jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _apply_bbox_augmentation(image, bbox, augmentation_func, *args):
    """Applies augmentation_func to the subsection of image indicated by bbox.

    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.
        augmentation_func: Augmentation function that will be applied to the
            subsection of image.
        *args: Additional parameters that will be passed into augmentation_func
            when it is called.

    Returns:
        A modified version of image, where the bbox location in the image will
        have `augmentation_func applied to it.
    """
    image_height = tf.cast(tf.shape(image)[0], tf.float32)
    image_width = tf.cast(tf.shape(image)[1], tf.float32)
    min_y = tf.to_int32(image_height * bbox[0])
    min_x = tf.to_int32(image_width * bbox[1])
    max_y = tf.to_int32(image_height * bbox[2])
    max_x = tf.to_int32(image_width * bbox[3])
    image_height = tf.to_int32(image_height)
    image_width = tf.to_int32(image_width)

    # Clip to be sure the max values do not fall out of range.
    max_y = tf.minimum(max_y, image_height - 1)
    max_x = tf.minimum(max_x, image_width - 1)

    # Get the sub-tensor that is the image within the bounding box region.
    bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

    # Apply the augmentation function to the bbox portion of the image.
    augmented_bbox_content = augmentation_func(bbox_content, *args)

    # Pad the augmented_bbox_content and the mask to match the shape of original
    # image.
    augmented_bbox_content = tf.pad(
        augmented_bbox_content,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]])

    # Create a mask that will be used to zero out a part of the original image.
    mask_tensor = tf.zeros_like(bbox_content)

    mask_tensor = tf.pad(
        mask_tensor,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]],
        constant_values=1)
    # Replace the old bbox content with the new augmented content.
    image = image * mask_tensor + augmented_bbox_content
    return image
//  [python]
def _shift_bbox(bbox, image_height, image_width, pixels, shift_horizontal):
    """Shifts the bbox coordinates by pixels.

    Args:
        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.
        image_height: Int, height of the image.
        image_width: Int, width of the image.
        pixels: An int. How many pixels to shift the bbox.
        shift_horizontal: Boolean. If true then shift in X dimension else shift in
            Y dimension.

    Returns:
        A tensor of the same shape as bbox, but now with the shifted coordinates.
    """
    pixels = tf.to_int32(pixels)
    # Convert bbox to integer pixel locations.
    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])

    if shift_horizontal:
        min_x = tf.maximum(0, min_x - pixels)
        max_x = tf.minimum(image_width, max_x - pixels)
    else:
        min_y = tf.maximum(0, min_y - pixels)
        max_y = tf.minimum(image_height, max_y - pixels)

    # Convert bbox back to floats.
    min_y = tf.to_float(min_y) / tf.to_float(image_height)
    min_x = tf.to_float(min_x) / tf.to_float(image_width)
    max_y = tf.to_float(max_y) / tf.to_float(image_height)
    max_x = tf.to_float(max_x) / tf.to_float(image_width)

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shift_bbox(bbox, image_height, image_width, pixels, shift_horizontal):
    """Shifts the bbox coordinates by pixels.

    Args:
        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.
        image_height: Int, height of the image.
        image_width: Int, width of the image.
        pixels: An int. How many pixels to shift the bbox.
        shift_horizontal: Boolean. If true then shift in X dimension else shift in
            Y dimension.

    Returns:
        A tensor of the same shape as bbox, but now with the shifted coordinates.
    """
    pixels = tf.cast(pixels, tf.int32)
    # Convert bbox to integer pixel locations.
    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)

    if shift_horizontal:
        min_x = tf.maximum(0, min_x - pixels)
        max_x = tf.minimum(image_width, max_x - pixels)
    else:
        min_y = tf.maximum(0, min_y - pixels)
        max_y = tf.minimum(image_height, max_y - pixels)

    # Convert bbox back to floats.
    min_y = tf.to_float(min_y) / tf.to_float(image_height)
    min_x = tf.to_float(min_x) / tf.to_float(image_width)
    max_y = tf.to_float(max_y) / tf.to_float(image_height)
    max_x = tf.to_float(max_x) / tf.to_float(image_width)

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [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 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 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 compute_losses(loss_dict, features, action_logits, obj_logits,
                   consumed_logits):
  """Compute the loss based on the logits and labels."""
  valid_obj_mask = tf.to_float(tf.greater(features["verbs"], 1))
  action_losses = tf.losses.sparse_softmax_cross_entropy(
      labels=features["verbs"],
      logits=action_logits,
      reduction=tf.losses.Reduction.NONE) * valid_obj_mask
  action_loss = tf.reduce_mean(action_losses)
  object_losses = tf.losses.sparse_softmax_cross_entropy(
      labels=features["objects"],
      logits=obj_logits,
      reduction=tf.losses.Reduction.NONE) * valid_obj_mask
  object_loss = tf.reduce_mean(object_losses)
  if "consumed" in features:
    consumed_loss = tf.reduce_mean(
        tf.losses.sparse_softmax_cross_entropy(
            labels=features["consumed"],
            logits=consumed_logits,
            reduction=tf.losses.Reduction.NONE) * valid_obj_mask)
  else:
    consumed_loss = 0.0
  loss_dict["grounding_loss"] = action_loss + object_loss + consumed_loss
  loss_dict["verbs_loss"] = action_loss
  loss_dict["objects_loss"] = object_loss
  loss_dict["verbs_losses"] = action_losses
  loss_dict["object_losses"] = object_losses
  loss_dict["consumed_loss"] = consumed_loss
  return loss_dict["grounding_loss"]
//  [python]
def compute_losses(loss_dict, features, action_logits, obj_logits,
                   consumed_logits):
  """Compute the loss based on the logits and labels."""
  valid_obj_mask = tf.cast(tf.greater(features["verbs"], 1), tf.float32)
  action_losses = tf.losses.sparse_softmax_cross_entropy(
      labels=features["verbs"],
      logits=action_logits,
      reduction=tf.losses.Reduction.NONE) * valid_obj_mask
  action_loss = tf.reduce_mean(action_losses)
  object_losses = tf.losses.sparse_softmax_cross_entropy(
      labels=features["objects"],
      logits=obj_logits,
      reduction=tf.losses.Reduction.NONE) * valid_obj_mask
  object_loss = tf.reduce_mean(object_losses)
  if "consumed" in features:
    consumed_loss = tf.reduce_mean(
        tf.losses.sparse_softmax_cross_entropy(
            labels=features["consumed"],
            logits=consumed_logits,
            reduction=tf.losses.Reduction.NONE) * valid_obj_mask)
  else:
    consumed_loss = 0.0
  loss_dict["grounding_loss"] = action_loss + object_loss + consumed_loss
  loss_dict["verbs_loss"] = action_loss
  loss_dict["objects_loss"] = object_loss
  loss_dict["verbs_losses"] = action_losses
  loss_dict["object_losses"] = object_losses
  loss_dict["consumed_loss"] = consumed_loss
  return loss_dict["grounding_loss"]
//  [python]
def _rotate_bbox(bbox, image_height, image_width, degrees):
    """Rotates the bbox coordinated by degrees.

    Args:
        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.
        image_height: Int, height of the image.
        image_width: Int, height of the image.
        degrees: Float, a scalar angle in degrees to rotate all images by. If
            degrees is positive the image will be rotated clockwise otherwise it will
            be rotated counterclockwise.

    Returns:
        A tensor of the same shape as bbox, but now with the rotated coordinates.
    """
    image_height, image_width = (
        tf.to_float(image_height), tf.to_float(image_width))

    # Convert from degrees to radians.
    degrees_to_radians = math.pi / 180.0
    radians = degrees * degrees_to_radians

    # Translate the bbox to the center of the image and turn the normalized 0-1
    # coordinates to absolute pixel locations.
    # Y coordinates are made negative as the y axis of images goes down with
    # increasing pixel values, so we negate to make sure x axis and y axis points
    # are in the traditionally positive direction.
    min_y = -tf.to_int32(image_height * (bbox[0] - 0.5))
    min_x = tf.to_int32(image_width * (bbox[1] - 0.5))
    max_y = -tf.to_int32(image_height * (bbox[2] - 0.5))
    max_x = tf.to_int32(image_width * (bbox[3] - 0.5))
    coordinates = tf.stack(
        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
    coordinates = tf.cast(coordinates, tf.float32)
    # Rotate the coordinates according to the rotation matrix clockwise if
    # radians is positive, else negative
    rotation_matrix = tf.stack(
        [[tf.cos(radians), tf.sin(radians)],
         [-tf.sin(radians), tf.cos(radians)]])
    new_coords = tf.cast(
        tf.matmul(rotation_matrix, tf.transpose(coordinates)), tf.int32)
    # Find min/max values and convert them back to normalized 0-1 floats.
    min_y = -(tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height - 0.5)
    min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width + 0.5
    max_y = -(tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height - 0.5)
    max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width + 0.5

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _rotate_bbox(bbox, image_height, image_width, degrees):
    """Rotates the bbox coordinated by degrees.

    Args:
        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.
        image_height: Int, height of the image.
        image_width: Int, height of the image.
        degrees: Float, a scalar angle in degrees to rotate all images by. If
            degrees is positive the image will be rotated clockwise otherwise it will
            be rotated counterclockwise.

    Returns:
        A tensor of the same shape as bbox, but now with the rotated coordinates.
    """
    image_height, image_width = (
        tf.to_float(image_height), tf.to_float(image_width))

    # Convert from degrees to radians.
    degrees_to_radians = math.pi / 180.0
    radians = degrees * degrees_to_radians

    # Translate the bbox to the center of the image and turn the normalized 0-1
    # coordinates to absolute pixel locations.
    # Y coordinates are made negative as the y axis of images goes down with
    # increasing pixel values, so we negate to make sure x axis and y axis points
    # are in the traditionally positive direction.
    min_y = -tf.cast(image_height * (bbox[0] - 0.5), tf.int32)
    min_x = tf.cast(image_width * (bbox[1] - 0.5), tf.int32)
    max_y = -tf.cast(image_height * (bbox[2] - 0.5), tf.int32)
    max_x = tf.cast(image_width * (bbox[3] - 0.5), tf.int32)
    coordinates = tf.stack(
        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
    coordinates = tf.cast(coordinates, tf.float32)
    # Rotate the coordinates according to the rotation matrix clockwise if
    # radians is positive, else negative
    rotation_matrix = tf.stack(
        [[tf.cos(radians), tf.sin(radians)],
         [-tf.sin(radians), tf.cos(radians)]])
    new_coords = tf.cast(
        tf.matmul(rotation_matrix, tf.transpose(coordinates)), tf.int32)
    # Find min/max values and convert them back to normalized 0-1 floats.
    min_y = -(tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height - 0.5)
    min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width + 0.5
    max_y = -(tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height - 0.5)
    max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width + 0.5

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
    """Shifts the bbox according to how the image was sheared.

    Args:
        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.
        image_height: Int, height of the image.
        image_width: Int, height of the image.
        level: Float. How much to shear the image.
        shear_horizontal: If true then shear in X dimension else shear in
            the Y dimension.

    Returns:
        A tensor of the same shape as bbox, but now with the shifted coordinates.
    """
    image_height, image_width = (
        tf.to_float(image_height), tf.to_float(image_width))

    # Change bbox coordinates to be pixels.
    min_y = tf.to_int32(image_height * bbox[0])
    min_x = tf.to_int32(image_width * bbox[1])
    max_y = tf.to_int32(image_height * bbox[2])
    max_x = tf.to_int32(image_width * bbox[3])
    coordinates = tf.stack(
        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
    coordinates = tf.cast(coordinates, tf.float32)

    # Shear the coordinates according to the translation matrix.
    if shear_horizontal:
        translation_matrix = tf.stack(
            [[1, 0], [-level, 1]])
    else:
        translation_matrix = tf.stack(
            [[1, -level], [0, 1]])
    translation_matrix = tf.cast(translation_matrix, tf.float32)
    new_coords = tf.cast(
        tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

    # Find min/max values and convert them back to floats.
    min_y = tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height
    min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width
    max_y = tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height
    max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
    """Shifts the bbox according to how the image was sheared.

    Args:
        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.
        image_height: Int, height of the image.
        image_width: Int, height of the image.
        level: Float. How much to shear the image.
        shear_horizontal: If true then shear in X dimension else shear in
            the Y dimension.

    Returns:
        A tensor of the same shape as bbox, but now with the shifted coordinates.
    """
    image_height, image_width = (
        tf.to_float(image_height), tf.to_float(image_width))

    # Change bbox coordinates to be pixels.
    min_y = tf.cast(image_height * bbox[0], tf.int32)
    min_x = tf.cast(image_width * bbox[1], tf.int32)
    max_y = tf.cast(image_height * bbox[2], tf.int32)
    max_x = tf.cast(image_width * bbox[3], tf.int32)
    coordinates = tf.stack(
        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
    coordinates = tf.cast(coordinates, tf.float32)

    # Shear the coordinates according to the translation matrix.
    if shear_horizontal:
        translation_matrix = tf.stack(
            [[1, 0], [-level, 1]])
    else:
        translation_matrix = tf.stack(
            [[1, -level], [0, 1]])
    translation_matrix = tf.cast(translation_matrix, tf.float32)
    new_coords = tf.cast(
        tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

    # Find min/max values and convert them back to floats.
    min_y = tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height
    min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width
    max_y = tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height
    max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def get_pred_mention_scores(self, candidate_starts,candidate_ends, pred_mention_starts, pred_mention_ends, pred_mention_scores):
    same_start = tf.equal(tf.expand_dims(pred_mention_starts, 1),
                          tf.expand_dims(candidate_starts, 0))  # [num_labeled, num_candidates]
    same_end = tf.equal(tf.expand_dims(pred_mention_ends, 1),
                        tf.expand_dims(candidate_ends, 0))  # [num_labeled, num_candidates]
    same_span = tf.logical_and(same_start, same_end)  # [num_labeled, num_candidates]
    candidate_scores = tf.matmul(tf.expand_dims(pred_mention_scores, 0), tf.to_float(same_span))  # [1, num_candidates]
    candidate_scores = tf.squeeze(candidate_scores, 0)  # [num_candidates]
    return candidate_scores
//  [python]
def get_pred_mention_scores(self, candidate_starts,candidate_ends, pred_mention_starts, pred_mention_ends, pred_mention_scores):
    same_start = tf.equal(tf.expand_dims(pred_mention_starts, 1),
                          tf.expand_dims(candidate_starts, 0))  # [num_labeled, num_candidates]
    same_end = tf.equal(tf.expand_dims(pred_mention_ends, 1),
                        tf.expand_dims(candidate_ends, 0))  # [num_labeled, num_candidates]
    same_span = tf.logical_and(same_start, same_end)  # [num_labeled, num_candidates]
    candidate_scores = tf.matmul(tf.expand_dims(pred_mention_scores, 0), tf.cast(same_span, tf.float32))  # [1, num_candidates]
    candidate_scores = tf.squeeze(candidate_scores, 0)  # [num_candidates]
    return candidate_scores
//  [python]
def action(self, context):
    """Samples rewards from posterior, and chooses best action accordingly."""

    # Round robin until each action has been selected "initial_pulls" times
    if self.t < self.hparams.num_actions * self.hparams.initial_pulls:
      return self.t % self.hparams.num_actions

    vals = []
    states, rewards, actions = bandit_utils.get_data_with_masked_rewards(
        self.data_h)
    state_action_pairs = tf.concat([states, actions], axis=-1)

    historical_x = tf.to_float(tf.expand_dims(state_action_pairs, axis=0))
    historical_y = tf.to_float(rewards.reshape(1, -1, 1))

    context = tf.to_float(context)
    tiled_context = tf.concat(
        [tf.tile(tf.reshape(context, [1, -1]), [self.hparams.num_actions, 1]),
         self._one_hot_vectors], axis=-1
    )
    target_x = tf.expand_dims(tiled_context, axis=0)
    target_y = None

    prediction = self.snp(historical_x, historical_y, target_x, target_y)
    vals = tf.squeeze(prediction.distribution.mean())

    return tf.argmax(vals).numpy()
//  [python]
def action(self, context):
    """Samples rewards from posterior, and chooses best action accordingly."""

    # Round robin until each action has been selected "initial_pulls" times
    if self.t < self.hparams.num_actions * self.hparams.initial_pulls:
      return self.t % self.hparams.num_actions

    vals = []
    states, rewards, actions = bandit_utils.get_data_with_masked_rewards(
        self.data_h)
    state_action_pairs = tf.concat([states, actions], axis=-1)

    historical_x = tf.cast(tf.expand_dims(state_action_pairs, axis=0), tf.float32)
    historical_y = tf.cast(rewards.reshape(1, -1, 1), tf.float32)

    context = tf.cast(context, tf.float32)
    tiled_context = tf.concat(
        [tf.tile(tf.reshape(context, [1, -1]), [self.hparams.num_actions, 1]),
         self._one_hot_vectors], axis=-1
    )
    target_x = tf.expand_dims(tiled_context, axis=0)
    target_y = None

    prediction = self.snp(historical_x, historical_y, target_x, target_y)
    vals = tf.squeeze(prediction.distribution.mean())

    return tf.argmax(vals).numpy()
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k)  # [k]
    antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
    antecedents_mask = antecedent_offsets >= 1  # [k, k]
    fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores,
                                                                                         0)  # [k, k]
    fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask))  # [k, k]
    fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)  # [k, k]

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False)  # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents)  # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents)  # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents)  # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k)  # [k]
    antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
    antecedents_mask = antecedent_offsets >= 1  # [k, k]
    fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores,
                                                                                         0)  # [k, k]
    fast_antecedent_scores += tf.log(tf.cast(antecedents_mask, dtype=tf.float32))  # [k, k]
    fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)  # [k, k]

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False)  # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents)  # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents)  # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents)  # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [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 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 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
//  [python]
def train(self, data, num_steps):
    """Trains the network for num_steps, using the provided data.

    Args:
      data: ContextualDataset object that provides the data.
      num_steps: Number of minibatches to train the network for.
    """

    if self.verbose:
      print('Training {} for {} steps...'.format(self.name, num_steps))

    avg_nll, avg_mse, avg_local_z_kl, avg_global_z_kl = 0., 0., 0., 0.
    for _ in range(num_steps):
      states, rewards, actions = bandit_utils.get_batch_with_masked_rewards(
          self.data_h,
          self.hparams.batch_size)
      state_action_pairs = tf.concat([states, actions], axis=-1)
      target_x = tf.to_float(tf.expand_dims(state_action_pairs, axis=0))
      target_y = tf.to_float(rewards.reshape(1, -1, 1))
      num_historical = tf.random.uniform(
          (), 1, max(2, target_x.shape[1]), dtype=tf.dtypes.int32)
      historical_x = target_x[:, :num_historical]
      historical_y = target_y[:, :num_historical]
      unseen_targets = target_y[:, num_historical:]
      data = (historical_x, historical_y, target_x, target_y, unseen_targets)

      nll, mse, local_z_kl, global_z_kl = self._step(self.snp,
                                                     data,
                                                     self._optimizer_config)
      avg_nll += nll
      avg_mse += mse
      avg_local_z_kl += local_z_kl
      avg_global_z_kl += global_z_kl
      self.gradient_updates.assign_add(1)

    if self.verbose:
      avg_nll /= num_steps
      avg_mse /= num_steps
      avg_local_z_kl /= num_steps
      avg_global_z_kl /= num_steps
      print('Average nll: {}, mse: {}, local kl: {} global kl: {}'
            .format(avg_nll, avg_mse, avg_local_z_kl, avg_global_z_kl))
//  [python]
def train(self, data, num_steps):
    """Trains the network for num_steps, using the provided data.

    Args:
      data: ContextualDataset object that provides the data.
      num_steps: Number of minibatches to train the network for.
    """

    if self.verbose:
      print('Training {} for {} steps...'.format(self.name, num_steps))

    avg_nll, avg_mse, avg_local_z_kl, avg_global_z_kl = 0., 0., 0., 0.
    for _ in range(num_steps):
      states, rewards, actions = bandit_utils.get_batch_with_masked_rewards(
          self.data_h,
          self.hparams.batch_size)
      state_action_pairs = tf.concat([states, actions], axis=-1)
      target_x = tf.cast(tf.expand_dims(state_action_pairs, axis=0), tf.float32)
      target_y = tf.cast(rewards.reshape(1, -1, 1), tf.float32)
      num_historical = tf.random.uniform(
          (), 1, max(2, target_x.shape[1]), dtype=tf.dtypes.int32)
      historical_x = target_x[:, :num_historical]
      historical_y = target_y[:, :num_historical]
      unseen_targets = target_y[:, num_historical:]
      data = (historical_x, historical_y, target_x, target_y, unseen_targets)

      nll, mse, local_z_kl, global_z_kl = self._step(self.snp,
                                                     data,
                                                     self._optimizer_config)
      avg_nll += nll
      avg_mse += mse
      avg_local_z_kl += local_z_kl
      avg_global_z_kl += global_z_kl
      self.gradient_updates.assign_add(1)

    if self.verbose:
      avg_nll /= num_steps
      avg_mse /= num_steps
      avg_local_z_kl /= num_steps
      avg_global_z_kl /= num_steps
      print('Average nll: {}, mse: {}, local kl: {} global kl: {}'
            .format(avg_nll, avg_mse, avg_local_z_kl, avg_global_z_kl))
//  [python]
def blend(image1, image2, factor):
    """Blend image1 and image2 using 'factor'.

    Factor can 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 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 infer(self,
            features=None,
            decode_length=50,
            beam_size=1,
            top_beams=1,
            alpha=0.0,
            use_tpu=False):
    """Produce predictions from the model."""
    if not features:
      features = {}
    inputs_old = None
    if "inputs" in features and len(features["inputs"].shape) < 4:
      inputs_old = features["inputs"]
      features["inputs"] = tf.expand_dims(features["inputs"], 2)

    # Create an initial targets tensor.
    if "partial_targets" in features:
      initial_output = tf.convert_to_tensor(features["partial_targets"])
    else:
      batch_size = common_layers.shape_list(features["inputs"])[0]
      length = common_layers.shape_list(features["inputs"])[1]
      target_length = tf.to_int32(2.0 * tf.to_float(length))
      initial_output = tf.zeros((batch_size, target_length, 1, 1),
                                dtype=tf.int64)

    features["targets"] = initial_output
    logits, _ = self(features)  # pylint: disable=not-callable
    samples = tf.argmax(logits, axis=-1)
    if inputs_old is not None:  # Restore to not confuse Estimator.
      features["inputs"] = inputs_old
    return samples
//  [python]
def infer(self,
            features=None,
            decode_length=50,
            beam_size=1,
            top_beams=1,
            alpha=0.0,
            use_tpu=False):
    """Produce predictions from the model."""
    if not features:
      features = {}
    inputs_old = None
    if "inputs" in features and len(features["inputs"].shape) < 4:
      inputs_old = features["inputs"]
      features["inputs"] = tf.expand_dims(features["inputs"], 2)

    # Create an initial targets tensor.
    if "partial_targets" in features:
      initial_output = tf.convert_to_tensor(features["partial_targets"])
    else:
      batch_size = common_layers.shape_list(features["inputs"])[0]
      length = common_layers.shape_list(features["inputs"])[1]
      target_length = tf.cast(2.0 * tf.cast(length, tf.float32), tf.int32)
      initial_output = tf.zeros((batch_size, target_length, 1, 1),
                                dtype=tf.int64)

    features["targets"] = initial_output
    logits, _ = self(features)  # pylint: disable=not-callable
    samples = tf.argmax(logits, axis=-1)
    if inputs_old is not None:  # Restore to not confuse Estimator.
      features["inputs"] = inputs_old
    return samples
//  [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 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))


        def scale_values(im):
            """Scale the image, making the lowest value 0 and the highest value 255."""
            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 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)


        def scale_values(im):
            """Scale the image, making the lowest value 0 and the highest value 255."""
            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 _string_to_bytes(text, max_length):
  """Given string and length, convert to byte seq of at most max_length.

  This process mimics docqa/elmo's preprocessing:
  https://github.com/allenai/document-qa/blob/master/docqa/elmo/data.py

  Note that we make use of BOS_CHAR_ID and EOS_CHAR_ID in iterator_utils.py &
  our usage differs from docqa/elmo.

  Args:
    text: tf.string tensor of shape []
    max_length: max number of chars for each word.

  Returns:
    A tf.int32 tensor of the byte encoded text.
  """
  byte_ids = tf.to_int32(tf.decode_raw(text, tf.uint8))
  byte_ids = byte_ids[:max_length - 2]
  padding = tf.fill([max_length - tf.shape(byte_ids)[0] - 2], PAD_CHAR_ID)
  byte_ids = tf.concat(
      [[BOW_CHAR_ID], byte_ids, [EOW_CHAR_ID], padding], axis=0)
  tf.logging.info(byte_ids)

  byte_ids = tf.reshape(byte_ids, [max_length])
  tf.logging.info(byte_ids.get_shape().as_list())
  return byte_ids + 1
//  [python]
def _string_to_bytes(text, max_length):
  """Given string and length, convert to byte seq of at most max_length.

  This process mimics docqa/elmo's preprocessing:
  https://github.com/allenai/document-qa/blob/master/docqa/elmo/data.py

  Note that we make use of BOS_CHAR_ID and EOS_CHAR_ID in iterator_utils.py &
  our usage differs from docqa/elmo.

  Args:
    text: tf.string tensor of shape []
    max_length: max number of chars for each word.

  Returns:
    A tf.int32 tensor of the byte encoded text.
  """
  byte_ids = tf.cast(tf.decode_raw(text, tf.uint8), tf.int32)
  byte_ids = byte_ids[:max_length - 2]
  padding = tf.fill([max_length - tf.shape(byte_ids)[0] - 2], PAD_CHAR_ID)
  byte_ids = tf.concat(
      [[BOW_CHAR_ID], byte_ids, [EOW_CHAR_ID], padding], axis=0)
  tf.logging.info(byte_ids)

  byte_ids = tf.reshape(byte_ids, [max_length])
  tf.logging.info(byte_ids.get_shape().as_list())
  return byte_ids + 1
//  [python]
def get_image_info(self):
    """Returns image information for scaled and original height and width."""
    return tf.stack([
        tf.to_float(self._scaled_height),
        tf.to_float(self._scaled_width),
        1.0 / self._image_scale,
        tf.to_float(self._ori_height),
        tf.to_float(self._ori_width)])
//  [python]
def get_image_info(self):
    """Returns image information for scaled and original height and width."""
    return tf.stack([
        tf.cast(self._scaled_height, tf.float32),
        tf.cast(self._scaled_width, tf.float32),
        1.0 / self._image_scale,
        tf.cast(self._ori_height, tf.float32),
        tf.cast(self._ori_width, tf.float32)])
//  [python]
def get_linear_loss_coef(loss_end_step):
  current_step = tf.train.get_or_create_global_step()
  return tf.to_float(current_step) / tf.to_float(loss_end_step)
//  [python]
def get_linear_loss_coef(loss_end_step):
  current_step = tf.train.get_or_create_global_step()
  return tf.cast(current_step, tf.float32) / tf.cast(loss_end_step, tf.float32)