message: Speeds up _random_flip for batched images.
diff: @@ -29,7 +29,6 @@ from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import control_flow_ops
  -from tensorflow.python.ops import functional_ops
  from tensorflow.python.ops import gen_image_ops
  from tensorflow.python.ops import gen_nn_ops
  from tensorflow.python.ops import math_ops
  @@ -301,11 +300,11 @@ def random_flip_left_right(image, seed=None):
  def _random_flip(image, flip_index, seed, scope_name):
  """Randomly (50% chance) flip an image along axis `flip_index`.
  +
  Args:
  image: 4-D Tensor of shape `[batch, height, width, channels]` or
  3-D Tensor of shape `[height, width, channels]`.
  - flip_index: The dimension along which to flip the image.
  - Vertical: 0, Horizontal: 1
  + flip_index: Dimension along which to flip image. Vertical: 0, Horizontal: 1
  seed: A Python integer. Used to create a random seed. See
  `tf.set_random_seed`
  for behavior.
  @@ -334,15 +333,16 @@ def _random_flip(image, flip_index, seed, scope_name):
  result = result[0] # TODO(b/111124878) remove this logic (CondV2).
  return fix_image_flip_shape(image, result)
  elif shape.ndims == 4:
  + batch_size = array_ops.shape(image)[0]
  uniform_random = random_ops.random_uniform(
  - [array_ops.shape(image)[0]], 0, 1.0, seed=seed
  + [batch_size], 0, 1.0, seed=seed
  )
  - mirror_cond = math_ops.less(uniform_random, .5)
  - return array_ops.where(
  - mirror_cond,
  - image,
  - functional_ops.map_fn(lambda x: array_ops.reverse(x, [flip_index]), image, dtype=image.dtype)
  + flips = math_ops.round(
  + array_ops.reshape(uniform_random, [batch_size, 1, 1, 1])
  )
  + flips = math_ops.cast(flips, image.dtype)
  + flipped_input = array_ops.reverse(image, [flip_index + 1])
  + return flips * flipped_input + (1 - flips) * image
  else:
  raise ValueError('\'image\' must have either 3 or 4 dimensions.')
  