def _model_with_preprocess_fn(model, model_preprocess_fn):
  """Combines the model with its preprocessing function.

  Args:
    model: Callable taking (preprocessed_images, is_training, test_local_stats)
      and returning logits.
    model_preprocess_fn: Image pre-processing to be combined with `model`.

  Returns:
    Callable taking (raw_images, is_training, test_local_stats) and returning
    logits.
  """

  def model_with_preprocess(images, **batchnorm_kwargs):
    # Transform inputs from [0, 1] to roughly [-1, +1], which is the range used
    # by ResNets. (This may entail slightly different transformations for each
    # channel to reflect the dataset statistics.)
    # Do so here (instead of during dataset loading) because
    # adversarial attacks assume that the model inputs are [0, 1].
    images = model_preprocess_fn(images)
    return model(images, **batchnorm_kwargs)

  return model_with_preprocess