jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.
  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    The cropped (and resized) image.
  Raises:
    ValueError: if `image` doesn't have rank of 3.
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  if len(image.get_shape().as_list()) != 3:
    raise ValueError('input must have rank of 3')
  original_channels = image.get_shape().as_list()[2]

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3), ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  image = tf.reshape(image, cropped_shape)
  image.set_shape([crop_height, crop_width, original_channels])
  return image
//  [python]
def train():
  """Train the model."""
  embeddings = load_embeddings()

  # Need a named parameter `param` since this will be called
  # with named arguments, so pylint: disable=unused-argument
  def model_function(features, labels, mode, params):
    """Builds the `tf.estimator.EstimatorSpec` to train/eval with."""
    is_train = mode == tf_estimator.ModeKeys.TRAIN
    logits = predict(is_train, embeddings, features["premise"],
                     features["hypothesis"])

    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=tf.to_int32(labels), logits=logits)
    loss = tf.reduce_mean(loss)
    if mode == tf_estimator.ModeKeys.TRAIN:
      train_op = get_train_op(loss)
    else:
      # Don't build the train_op unnecessarily, since the ADAM variables can
      # cause problems with loading checkpoints on CPUs.
      train_op = None
    metrics = dict(
        accuracy=tf.metrics.accuracy(
            tf.argmax(logits, 1, output_type=tf.int32), tf.to_int32(labels)))

    checkpoint_file = FLAGS.checkpoint_file
    if checkpoint_file is None:
      scaffold = None
    else:
      saver = tf.train.Saver(tf.trainable_variables())

      def _init_fn(_, sess):
        saver.restore(sess, checkpoint_file)

      scaffold = tf.train.Scaffold(init_fn=_init_fn)

    return tf_estimator.EstimatorSpec(
        mode=mode,
        scaffold=scaffold,
        loss=loss,
        predictions=None,
        train_op=train_op,
        eval_metric_ops=metrics)

  def compare_fn(best_eval_result, current_eval_result):
    return best_eval_result["accuracy"] < current_eval_result["accuracy"]

  exporter = best_checkpoint_exporter.BestCheckpointExporter(
      event_file_pattern="eval_default/*.tfevents.*",
      compare_fn=compare_fn,
  )

  experiment_utils.run_experiment(
      model_fn=model_function,
      train_input_fn=lambda: load_batched_dataset(True, embeddings),
      eval_input_fn=lambda: load_batched_dataset(False, embeddings),
      exporters=[exporter])
//  [python]
def train():
  """Train the model."""
  embeddings = load_embeddings()

  # Need a named parameter `param` since this will be called
  # with named arguments, so pylint: disable=unused-argument
  def model_function(features, labels, mode, params):
    """Builds the `tf.estimator.EstimatorSpec` to train/eval with."""
    is_train = mode == tf_estimator.ModeKeys.TRAIN
    logits = predict(is_train, embeddings, features["premise"],
                     features["hypothesis"])

    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=tf.cast(labels, tf.int32), logits=logits)
    loss = tf.reduce_mean(loss)
    if mode == tf_estimator.ModeKeys.TRAIN:
      train_op = get_train_op(loss)
    else:
      # Don't build the train_op unnecessarily, since the ADAM variables can
      # cause problems with loading checkpoints on CPUs.
      train_op = None
    metrics = dict(
        accuracy=tf.metrics.accuracy(
            tf.argmax(logits, 1, output_type=tf.int32), tf.cast(labels, tf.int32)))

    checkpoint_file = FLAGS.checkpoint_file
    if checkpoint_file is None:
      scaffold = None
    else:
      saver = tf.train.Saver(tf.trainable_variables())

      def _init_fn(_, sess):
        saver.restore(sess, checkpoint_file)

      scaffold = tf.train.Scaffold(init_fn=_init_fn)

    return tf_estimator.EstimatorSpec(
        mode=mode,
        scaffold=scaffold,
        loss=loss,
        predictions=None,
        train_op=train_op,
        eval_metric_ops=metrics)

  def compare_fn(best_eval_result, current_eval_result):
    return best_eval_result["accuracy"] < current_eval_result["accuracy"]

  exporter = best_checkpoint_exporter.BestCheckpointExporter(
      event_file_pattern="eval_default/*.tfevents.*",
      compare_fn=compare_fn,
  )

  experiment_utils.run_experiment(
      model_fn=model_function,
      train_input_fn=lambda: load_batched_dataset(True, embeddings),
      eval_input_fn=lambda: load_batched_dataset(False, embeddings),
      exporters=[exporter])
//  [python]
def evaluate():
  """Evaluate a model on the dev set."""
  sess = tf.Session()
  tf.logging.info("Building graph...")

  embeddings = load_embeddings()
  tf_data = load_batched_dataset(False, embeddings)
  it = tf_data.make_initializable_iterator()
  features, labels = it.get_next()

  logits = predict(False, embeddings, features["premise"],
                   features["hypothesis"])
  accuracy, update_ops = tf.metrics.accuracy(
      tf.argmax(logits, 1, output_type=tf.int32), tf.to_int32(labels))

  tf.logging.info("Running initializers...")
  checkpoint_file = FLAGS.checkpoint_file
  if checkpoint_file is not None:
    saver = tf.train.Saver(tf.trainable_variables())
    tf.logging.info("Restoring from checkpoint: " + checkpoint_file)
    saver.restore(sess, checkpoint_file)
  else:
    tf.logging.warning("No checkpoint given, evaling model with random weights")
    sess.run(tf.global_variables_initializer())
  sess.run(tf.local_variables_initializer())
  sess.run(tf.tables_initializer())
  sess.run(it.initializer)

  tf.logging.info("Starting loop....")
  while True:
    try:
      sess.run(update_ops)
    except tf.errors.OutOfRangeError:
      break
  tf.logging.info("Done")

  accuracy = sess.run(accuracy)
  print("Accuracy: %f" % accuracy)
//  [python]
def evaluate():
  """Evaluate a model on the dev set."""
  sess = tf.Session()
  tf.logging.info("Building graph...")

  embeddings = load_embeddings()
  tf_data = load_batched_dataset(False, embeddings)
  it = tf_data.make_initializable_iterator()
  features, labels = it.get_next()

  logits = predict(False, embeddings, features["premise"],
                   features["hypothesis"])
  accuracy, update_ops = tf.metrics.accuracy(
      tf.argmax(logits, 1, output_type=tf.int32), tf.cast(labels, tf.int32))

  tf.logging.info("Running initializers...")
  checkpoint_file = FLAGS.checkpoint_file
  if checkpoint_file is not None:
    saver = tf.train.Saver(tf.trainable_variables())
    tf.logging.info("Restoring from checkpoint: " + checkpoint_file)
    saver.restore(sess, checkpoint_file)
  else:
    tf.logging.warning("No checkpoint given, evaling model with random weights")
    sess.run(tf.global_variables_initializer())
  sess.run(tf.local_variables_initializer())
  sess.run(tf.tables_initializer())
  sess.run(it.initializer)

  tf.logging.info("Starting loop....")
  while True:
    try:
      sess.run(update_ops)
    except tf.errors.OutOfRangeError:
      break
  tf.logging.info("Done")

  accuracy = sess.run(accuracy)
  print("Accuracy: %f" % accuracy)
//  [python]
def preprocess_for_eval(image,
                        output_height,
                        output_width,
                        add_image_summaries=True,
                        use_grayscale=False):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))
  # Transform the image to floats.
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)

  # Resize and crop if needed.
  resized_image = tf.image.resize_image_with_crop_or_pad(image,
                                                         output_width,
                                                         output_height)
  if add_image_summaries:
    tf.summary.image('resized_image', tf.expand_dims(resized_image, 0))

  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(resized_image)
//  [python]
def preprocess_for_eval(image,
                        output_height,
                        output_width,
                        add_image_summaries=True,
                        use_grayscale=False):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))
  # Transform the image to floats.
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)

  # Resize and crop if needed.
  resized_image = tf.image.resize_image_with_crop_or_pad(image,
                                                         output_width,
                                                         output_height)
  if add_image_summaries:
    tf.summary.image('resized_image', tf.expand_dims(resized_image, 0))

  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(resized_image)
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can_scripts be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.

  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.to_float(image1)
  image2 = tf.to_float(image2)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.to_float(image1) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can_scripts be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.

  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.cast(image1, tf.float32)
  image2 = tf.cast(image2, tf.float32)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.cast(image1, tf.float32) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "probs": tf.FixedLenFeature([2], tf.float32)
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "probs": tf.FixedLenFeature([2], tf.float32)
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         padding=_PADDING,
                         add_image_summaries=True,
                         use_grayscale=False):
  """Preprocesses the given image for training.

  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    padding: The amound of padding before and after each dimension of the image.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))

  # Transform the image to floats.
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  if padding > 0:
    image = tf.pad(image, [[padding, padding], [padding, padding], [0, 0]])
  # Randomly crop a [height, width] section of the image.
  distorted_image = tf.random_crop(image,
                                   [output_height, output_width, 3])

  # Randomly flip the image horizontally.
  distorted_image = tf.image.random_flip_left_right(distorted_image)

  if add_image_summaries:
    tf.summary.image('distorted_image', tf.expand_dims(distorted_image, 0))

  # Because these operations are not commutative, consider randomizing
  # the order their operation.
  distorted_image = tf.image.random_brightness(distorted_image,
                                               max_delta=63)
  distorted_image = tf.image.random_contrast(distorted_image,
                                             lower=0.2, upper=1.8)
  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(distorted_image)
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         padding=_PADDING,
                         add_image_summaries=True,
                         use_grayscale=False):
  """Preprocesses the given image for training.

  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    padding: The amound of padding before and after each dimension of the image.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))

  # Transform the image to floats.
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  if padding > 0:
    image = tf.pad(image, [[padding, padding], [padding, padding], [0, 0]])
  # Randomly crop a [height, width] section of the image.
  distorted_image = tf.random_crop(image,
                                   [output_height, output_width, 3])

  # Randomly flip the image horizontally.
  distorted_image = tf.image.random_flip_left_right(distorted_image)

  if add_image_summaries:
    tf.summary.image('distorted_image', tf.expand_dims(distorted_image, 0))

  # Because these operations are not commutative, consider randomizing
  # the order their operation.
  distorted_image = tf.image.random_brightness(distorted_image,
                                               max_delta=63)
  distorted_image = tf.image.random_contrast(distorted_image,
                                             lower=0.2, upper=1.8)
  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(distorted_image)
//  [python]
def recall_at_k(labels, embeddings, k, embed_normed=True):
    # make sure embedding should be l2-normalized
    if not embed_normed:
        embeddings = tf.nn.l2_normalize(embeddings, axis=1)
    # batch_size = tf.size(labels)
    embed_shape = modeling.get_shape_list(embeddings)
    batch_size = embed_shape[0]

    sim_mat = tf.matmul(embeddings, embeddings, transpose_b=True)
    sim_mat = sim_mat - tf.eye(batch_size) * 2.0

    labels = tf.expand_dims(labels, -1)
    mask = tf.equal(labels, tf.transpose(labels))  # shape: (batch_size, batch_size)
    eye = tf.eye(batch_size, dtype=tf.bool)
    mask = tf.logical_and(mask, tf.logical_not(eye))
    mask_pos = tf.where(mask, sim_mat, -tf.ones_like(sim_mat))  # shape: (batch_size, batch_size)

    if isinstance(k, int):
        _, pos_top_k_idx = tf.nn.top_k(mask_pos, k)  # shape: (batch_size, k)
        return tf.metrics.recall_at_k(labels=tf.to_int64(pos_top_k_idx), predictions=sim_mat, k=k)
    if any((isinstance(k, list), isinstance(k, tuple), isinstance(k, set))):
        metrics = {}
        for kk in k:
            if k < 1:
                continue
            _, pos_top_k_idx = tf.nn.top_k(mask_pos, kk)
            metrics["recall@"+str(kk)] = tf.metrics.recall_at_k(labels=tf.to_int64(pos_top_k_idx), predictions=sim_mat, k=kk)
        return metrics
    raise ValueError("k should be a `int` or a list/tuple/set of int.")
//  [python]
def recall_at_k(labels, embeddings, k, embed_normed=True):
    # make sure embedding should be l2-normalized
    if not embed_normed:
        embeddings = tf.nn.l2_normalize(embeddings, axis=1)
    # batch_size = tf.size(labels)
    embed_shape = modeling.get_shape_list(embeddings)
    batch_size = embed_shape[0]

    sim_mat = tf.matmul(embeddings, embeddings, transpose_b=True)
    sim_mat = sim_mat - tf.eye(batch_size) * 2.0

    labels = tf.expand_dims(labels, -1)
    mask = tf.equal(labels, tf.transpose(labels))  # shape: (batch_size, batch_size)
    eye = tf.eye(batch_size, dtype=tf.bool)
    mask = tf.logical_and(mask, tf.logical_not(eye))
    mask_pos = tf.where(mask, sim_mat, -tf.ones_like(sim_mat))  # shape: (batch_size, batch_size)

    if isinstance(k, int):
        _, pos_top_k_idx = tf.nn.top_k(mask_pos, k)  # shape: (batch_size, k)
        return tf.metrics.recall_at_k(labels=tf.cast(pos_top_k_idx, tf.int64), predictions=sim_mat, k=k)
    if any((isinstance(k, list), isinstance(k, tuple), isinstance(k, set))):
        metrics = {}
        for kk in k:
            if k < 1:
                continue
            _, pos_top_k_idx = tf.nn.top_k(mask_pos, kk)
            metrics["recall@"+str(kk)] = tf.metrics.recall_at_k(labels=tf.cast(pos_top_k_idx, tf.int64), predictions=sim_mat, k=kk)
        return metrics
    raise ValueError("k should be a `int` or a list/tuple/set of int.")
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can_scripts be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.to_float(tf.reduce_min(image))
    hi = tf.to_float(tf.reduce_max(image))

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.to_float(im) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can_scripts be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.cast(tf.reduce_min(image), tf.float32)
    hi = tf.cast(tf.reduce_max(image), tf.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.cast(im, tf.float32) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def padded_accuracy(logits, labels):
  """Percentage of times that predictions matches labels on non-0s."""
  with tf.variable_scope("padded_accuracy", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.to_float(tf.not_equal(labels, 0))
    outputs = tf.to_int32(tf.argmax(logits, axis=-1))
    padded_labels = tf.to_int32(labels)
    return tf.to_float(tf.equal(outputs, padded_labels)), weights
//  [python]
def padded_accuracy(logits, labels):
  """Percentage of times that predictions matches labels on non-0s."""
  with tf.variable_scope("padded_accuracy", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    outputs = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
    padded_labels = tf.cast(labels, tf.int32)
    return tf.cast(tf.equal(outputs, padded_labels), tf.float32), weights
//  [python]
def rouge_2_fscore(logits, labels):
  """ROUGE-2 F1 score computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

  Args:
    logits: tensor, model predictions
    labels: tensor, gold output.

  Returns:
    rouge2_fscore: approx rouge-2 f1 score.
  """
  predictions = tf.to_int32(tf.argmax(logits, axis=-1))
  # TODO: Look into removing use of py_func
  rouge_2_f_score = tf.py_func(rouge_n, (predictions, labels), tf.float32)
  return rouge_2_f_score, tf.constant(1.0)
//  [python]
def rouge_2_fscore(logits, labels):
  """ROUGE-2 F1 score computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

  Args:
    logits: tensor, model predictions
    labels: tensor, gold output.

  Returns:
    rouge2_fscore: approx rouge-2 f1 score.
  """
  predictions = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
  # TODO: Look into removing use of py_func
  rouge_2_f_score = tf.py_func(rouge_n, (predictions, labels), tf.float32)
  return rouge_2_f_score, tf.constant(1.0)
//  [python]
def padded_accuracy_topk(logits, labels, k):
  """Percentage of times that top-k predictions matches labels on non-0s."""
  with tf.variable_scope("padded_accuracy_topk", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.to_float(tf.not_equal(labels, 0))
    effective_k = tf.minimum(k, tf.shape(logits)[-1])
    _, outputs = tf.nn.top_k(logits, k=effective_k)
    outputs = tf.to_int32(outputs)
    padded_labels = tf.to_int32(labels)
    padded_labels = tf.expand_dims(padded_labels, axis=-1)
    padded_labels += tf.zeros_like(outputs)  # Pad to same shape.
    same = tf.to_float(tf.equal(outputs, padded_labels))
    same_topk = tf.reduce_sum(same, axis=-1)
    return same_topk, weights
//  [python]
def padded_accuracy_topk(logits, labels, k):
  """Percentage of times that top-k predictions matches labels on non-0s."""
  with tf.variable_scope("padded_accuracy_topk", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    effective_k = tf.minimum(k, tf.shape(logits)[-1])
    _, outputs = tf.nn.top_k(logits, k=effective_k)
    outputs = tf.cast(outputs, tf.int32)
    padded_labels = tf.cast(labels, tf.int32)
    padded_labels = tf.expand_dims(padded_labels, axis=-1)
    padded_labels += tf.zeros_like(outputs)  # Pad to same shape.
    same = tf.cast(tf.equal(outputs, padded_labels), tf.float32)
    same_topk = tf.reduce_sum(same, axis=-1)
    return same_topk, weights
//  [python]
def bleu_score(logits, labels):
  """Approximate BLEU score computation between labels and predictions.

  An approximate BLEU scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4
  and use brevity penalty. Also, this does not have beam search.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch-size, length_labels]

  Returns:
    bleu: int, approx bleu score
  """
  predictions = tf.to_int32(tf.argmax(logits, axis=-1))
  # TODO: Look into removing use of py_func
  bleu = tf.py_func(compute_bleu, (labels, predictions), tf.float32)
  return bleu, tf.constant(1.0)
//  [python]
def bleu_score(logits, labels):
  """Approximate BLEU score computation between labels and predictions.

  An approximate BLEU scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4
  and use brevity penalty. Also, this does not have beam search.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch-size, length_labels]

  Returns:
    bleu: int, approx bleu score
  """
  predictions = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
  # TODO: Look into removing use of py_func
  bleu = tf.py_func(compute_bleu, (labels, predictions), tf.float32)
  return bleu, tf.constant(1.0)
//  [python]
def padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary
  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  with tf.name_scope("loss", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)

    # Calculate smoothing cross entropy
    with tf.name_scope("smoothing_cross_entropy", values=[logits, labels]):
      confidence = 1.0 - smoothing
      low_confidence = (1.0 - confidence) / tf.to_float(vocab_size - 1)
      soft_targets = tf.one_hot(
          tf.cast(labels, tf.int32),
          depth=vocab_size,
          on_value=confidence,
          off_value=low_confidence)
      xentropy = tf.nn.softmax_cross_entropy_with_logits_v2(
          logits=logits, labels=soft_targets)

      # Calculate the best (lowest) possible value of cross entropy, and
      # subtract from the cross entropy loss.
      normalizing_constant = -(
          confidence * tf.log(confidence) + tf.to_float(vocab_size - 1) *
          low_confidence * tf.log(low_confidence + 1e-20))
      xentropy -= normalizing_constant

    weights = tf.to_float(tf.not_equal(labels, 0))
    return xentropy * weights, weights
//  [python]
def padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary
  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  with tf.name_scope("loss", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)

    # Calculate smoothing cross entropy
    with tf.name_scope("smoothing_cross_entropy", values=[logits, labels]):
      confidence = 1.0 - smoothing
      low_confidence = (1.0 - confidence) / tf.cast(vocab_size - 1, tf.float32)
      soft_targets = tf.one_hot(
          tf.cast(labels, tf.int32),
          depth=vocab_size,
          on_value=confidence,
          off_value=low_confidence)
      xentropy = tf.nn.softmax_cross_entropy_with_logits_v2(
          logits=logits, labels=soft_targets)

      # Calculate the best (lowest) possible value of cross entropy, and
      # subtract from the cross entropy loss.
      normalizing_constant = -(
          confidence * tf.log(confidence) + tf.cast(vocab_size - 1, tf.float32) *
          low_confidence * tf.log(low_confidence + 1e-20))
      xentropy -= normalizing_constant

    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    return xentropy * weights, weights
//  [python]
def rouge_l_fscore(predictions, labels):
  """ROUGE scores computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.

  Returns:
    rouge_l_fscore: approx rouge-l f1 score.
  """
  outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
  rouge_l_f_score = tf.py_func(rouge_l_sentence_level, (outputs, labels),
                               tf.float32)
  return rouge_l_f_score, tf.constant(1.0)
//  [python]
def rouge_l_fscore(predictions, labels):
  """ROUGE scores computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.

  Returns:
    rouge_l_fscore: approx rouge-l f1 score.
  """
  outputs = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
  rouge_l_f_score = tf.py_func(rouge_l_sentence_level, (outputs, labels),
                               tf.float32)
  return rouge_l_f_score, tf.constant(1.0)
//  [python]
def padded_sequence_accuracy(logits, labels):
  """Percentage of times that predictions matches labels everywhere (non-0)."""
  with tf.variable_scope("padded_sequence_accuracy", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.to_float(tf.not_equal(labels, 0))
    outputs = tf.to_int32(tf.argmax(logits, axis=-1))
    padded_labels = tf.to_int32(labels)
    not_correct = tf.to_float(tf.not_equal(outputs, padded_labels)) * weights
    axis = list(range(1, len(outputs.get_shape())))
    correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
    return correct_seq, tf.constant(1.0)
//  [python]
def padded_sequence_accuracy(logits, labels):
  """Percentage of times that predictions matches labels everywhere (non-0)."""
  with tf.variable_scope("padded_sequence_accuracy", values=[logits, labels]):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    outputs = tf.to_int32(tf.argmax(logits, axis=-1))
    padded_labels = tf.to_int32(labels)
    not_correct = tf.cast(tf.not_equal(outputs, padded_labels), tf.float32) * weights
    axis = list(range(1, len(outputs.get_shape())))
    correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
    return correct_seq, tf.constant(1.0)
//  [python]
def potential_rewards(states,
                      actions,
                      rewards,
                      next_states,
                      contexts,
                      gamma=1.0,
                      reward_fn=None):
  """Return the potential-based rewards.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    gamma: Reward discount.
    reward_fn: A reward function.
  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del actions  # unused args
  gamma = tf.to_float(gamma)
  rewards_tp1, discounts = reward_fn(None, None, rewards, next_states, contexts)
  rewards, _ = reward_fn(None, None, rewards, states, contexts)
  return -rewards + gamma * rewards_tp1, discounts
//  [python]
def potential_rewards(states,
                      actions,
                      rewards,
                      next_states,
                      contexts,
                      gamma=1.0,
                      reward_fn=None):
  """Return the potential-based rewards.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    gamma: Reward discount.
    reward_fn: A reward function.
  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del actions  # unused args
  gamma = tf.cast(gamma, tf.float32)
  rewards_tp1, discounts = reward_fn(None, None, rewards, next_states, contexts)
  rewards, _ = reward_fn(None, None, rewards, states, contexts)
  return -rewards + gamma * rewards_tp1, discounts
//  [python]
def diff_rewards(
    states,
    actions,
    rewards,
    next_states,
    contexts,
    state_indices=None,
    goal_index=0,):
  """Returns (next_states - goals) as a batched vector reward."""
  del states, rewards, actions  # Unused
  if state_indices is not None:
    next_states = index_states(next_states, state_indices)
  rewards = tf.to_float(next_states - contexts[goal_index])
  return rewards, tf.ones_like(rewards)
//  [python]
def diff_rewards(
    states,
    actions,
    rewards,
    next_states,
    contexts,
    state_indices=None,
    goal_index=0,):
  """Returns (next_states - goals) as a batched vector reward."""
  del states, rewards, actions  # Unused
  if state_indices is not None:
    next_states = index_states(next_states, state_indices)
  rewards = tf.cast(next_states - contexts[goal_index], tf.float32)
  return rewards, tf.ones_like(rewards)
//  [python]
def ctrl_rewards(states,
                 actions,
                 rewards,
                 next_states,
                 contexts,
                 reward_scales=1.0):
  """Returns the negative control cost.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
      must be broadcastable to number of reward dimensions.

  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del states, rewards, contexts  # Unused
  if actions is None:
    rewards = tf.to_float(tf.zeros(shape=next_states.shape[:1]))
  else:
    rewards = -tf.reduce_sum(tf.square(actions), axis=1)
    rewards *= reward_scales
    rewards = tf.to_float(rewards)
  return rewards, tf.ones_like(rewards)
//  [python]
def ctrl_rewards(states,
                 actions,
                 rewards,
                 next_states,
                 contexts,
                 reward_scales=1.0):
  """Returns the negative control cost.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
      must be broadcastable to number of reward dimensions.

  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del states, rewards, contexts  # Unused
  if actions is None:
    rewards = tf.cast(tf.zeros(shape=next_states.shape[:1]), dtype=tf.float32)
  else:
    rewards = -tf.reduce_sum(tf.square(actions), axis=1)
    rewards *= reward_scales
    rewards = tf.cast(rewards, dtype=tf.float32)
  return rewards, tf.ones_like(rewards)
//  [python]
def _cutout_inside_bbox(image, bbox, pad_fraction):
  """Generates cutout mask and the mean pixel value of the bbox.

  First a location is randomly chosen within the image as the center where the
  cutout mask will be applied. Note this can_scripts be towards the boundaries of the
  image, so the full cutout mask may not be applied.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    pad_fraction: Float that specifies how large the cutout mask should be in
      in reference to the size of the original bbox. If pad_fraction is 0.25,
      then the cutout mask will be of shape
      (0.25 * bbox height, 0.25 * bbox width).

  Returns:
    A tuple. Fist element is a tensor of the same shape as image where each
    element is either a 1 or 0 that is used to determine where the image
    will have cutout applied. The second element is the mean of the pixels
    in the image where the bbox is located.
  """
  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  # Transform from shape [1, 4] to [4].
  bbox = tf.squeeze(bbox)

  min_y = tf.to_int32(tf.to_float(image_height) * bbox[0])
  min_x = tf.to_int32(tf.to_float(image_width) * bbox[1])
  max_y = tf.to_int32(tf.to_float(image_height) * bbox[2])
  max_x = tf.to_int32(tf.to_float(image_width) * bbox[3])

  # Calculate the mean pixel values in the bounding box, which will be used
  # to fill the cutout region.
  mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                        reduction_indices=[0, 1])

  # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
  # region lies entirely within the bbox.
  box_height = max_y - min_y + 1
  box_width = max_x - min_x + 1
  pad_size_height = tf.to_int32(pad_fraction * (box_height / 2))
  pad_size_width = tf.to_int32(pad_fraction * (box_width / 2))

  # Sample the center location in the image where the zero mask will be applied.
  cutout_center_height = tf.random_uniform(
      shape=[], minval=min_y, maxval=max_y+1,
      dtype=tf.int32)

  cutout_center_width = tf.random_uniform(
      shape=[], minval=min_x, maxval=max_x+1,
      dtype=tf.int32)

  lower_pad = tf.maximum(
      0, cutout_center_height - pad_size_height)
  upper_pad = tf.maximum(
      0, image_height - cutout_center_height - pad_size_height)
  left_pad = tf.maximum(
      0, cutout_center_width - pad_size_width)
  right_pad = tf.maximum(
      0, image_width - cutout_center_width - pad_size_width)

  cutout_shape = [image_height - (lower_pad + upper_pad),
                  image_width - (left_pad + right_pad)]
  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

  mask = tf.pad(
      tf.zeros(cutout_shape, dtype=image.dtype),
      padding_dims, constant_values=1)

  mask = tf.expand_dims(mask, 2)
  mask = tf.tile(mask, [1, 1, 3])

  return mask, mean
//  [python]
def _cutout_inside_bbox(image, bbox, pad_fraction):
  """Generates cutout mask and the mean pixel value of the bbox.

  First a location is randomly chosen within the image as the center where the
  cutout mask will be applied. Note this can_scripts be towards the boundaries of the
  image, so the full cutout mask may not be applied.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    pad_fraction: Float that specifies how large the cutout mask should be in
      in reference to the size of the original bbox. If pad_fraction is 0.25,
      then the cutout mask will be of shape
      (0.25 * bbox height, 0.25 * bbox width).

  Returns:
    A tuple. Fist element is a tensor of the same shape as image where each
    element is either a 1 or 0 that is used to determine where the image
    will have cutout applied. The second element is the mean of the pixels
    in the image where the bbox is located.
  """
  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  # Transform from shape [1, 4] to [4].
  bbox = tf.squeeze(bbox)

  min_y = tf.cast(tf.to_float(image_height) * bbox[0], tf.int32)
  min_x = tf.cast(tf.to_float(image_width) * bbox[1], tf.int32)
  max_y = tf.cast(tf.to_float(image_height) * bbox[2], tf.int32)
  max_x = tf.cast(tf.to_float(image_width) * bbox[3], tf.int32)

  # Calculate the mean pixel values in the bounding box, which will be used
  # to fill the cutout region.
  mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                        reduction_indices=[0, 1])

  # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
  # region lies entirely within the bbox.
  box_height = max_y - min_y + 1
  box_width = max_x - min_x + 1
  pad_size_height = tf.cast(pad_fraction * (box_height / 2), tf.int32)
  pad_size_width = tf.cast(pad_fraction * (box_width / 2), tf.int32)

  # Sample the center location in the image where the zero mask will be applied.
  cutout_center_height = tf.random_uniform(
      shape=[], minval=min_y, maxval=max_y+1,
      dtype=tf.int32)

  cutout_center_width = tf.random_uniform(
      shape=[], minval=min_x, maxval=max_x+1,
      dtype=tf.int32)

  lower_pad = tf.maximum(
      0, cutout_center_height - pad_size_height)
  upper_pad = tf.maximum(
      0, image_height - cutout_center_height - pad_size_height)
  left_pad = tf.maximum(
      0, cutout_center_width - pad_size_width)
  right_pad = tf.maximum(
      0, image_width - cutout_center_width - pad_size_width)

  cutout_shape = [image_height - (lower_pad + upper_pad),
                  image_width - (left_pad + right_pad)]
  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

  mask = tf.pad(
      tf.zeros(cutout_shape, dtype=image.dtype),
      padding_dims, constant_values=1)

  mask = tf.expand_dims(mask, 2)
  mask = tf.tile(mask, [1, 1, 3])

  return mask, mean