def wrap(image):
  """Returns 'image' with an extra channel set to all 1s."""
  shape = tf.shape(image)
  extended_channel = tf.ones([shape[0], shape[1], 1], image.dtype)
  extended = tf.concat([image, extended_channel], 2)
  return extended
Returns 'image' with an extra channel set to all 1s.
inference_api/src/main/object_detection/utils/autoaugment_utils.py