jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def input_fn_builder(input_files, seq_length, qry_length, is_training,
                     drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "doc_input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "doc_input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "doc_segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "qry_input_ids": tf.FixedLenFeature([qry_length], tf.int64),
      "qry_input_mask": tf.FixedLenFeature([qry_length], tf.int64),
      "qry_segment_ids": tf.FixedLenFeature([qry_length], tf.int64),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)

  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.Dataset.from_tensor_slices(input_files).interleave(
        lambda x: tf.data.TFRecordDataset(x),  # pylint: disable=unnecessary-lambda
        cycle_length=len(input_files),
        block_length=1)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.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 input_fn_builder(input_files, seq_length, qry_length, is_training,
                     drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "doc_input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "doc_input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "doc_segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "qry_input_ids": tf.FixedLenFeature([qry_length], tf.int64),
      "qry_input_mask": tf.FixedLenFeature([qry_length], tf.int64),
      "qry_segment_ids": tf.FixedLenFeature([qry_length], tf.int64),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)

  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.Dataset.from_tensor_slices(input_files).interleave(
        lambda x: tf.data.TFRecordDataset(x),  # pylint: disable=unnecessary-lambda
        cycle_length=len(input_files),
        block_length=1)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.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 __call__(self, features, is_training):
        seqs_id = features['seqs_i']

        # positional encoding
        seqs_units = self.item_embs(seqs_id)
        seqs_units = self.pcoding(seqs_units)

        # Dropout
        with tf.variable_scope("input_transform"):
            seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate,
                                           training=tf.convert_to_tensor(is_training))
            seqs_masks = tf.expand_dims(tf.to_float(tf.not_equal(seqs_id, 0)), -1)

        # multi-head attention
        seqs_outs = seqs_units * seqs_masks
        for i, (attention, dense) in enumerate(zip(self.list_attention, self.list_dense)):
            with tf.variable_scope("num_blocks_%d" % i):
                with tf.variable_scope("attention"):
                    seqs_outs = attention(layernorm(seqs_outs), seqs_outs, is_training, causality=True)

                with tf.variable_scope("feedforward"):
                    seqs_outs = dense(layernorm(seqs_outs), is_training)
                    seqs_outs *= seqs_masks

        with tf.variable_scope("output_ln"):
            seqs_outs = layernorm(seqs_outs)

        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']

        # positional encoding
        seqs_units = self.item_embs(seqs_id)
        seqs_units = self.pcoding(seqs_units)

        # Dropout
        with tf.variable_scope("input_transform"):
            seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate,
                                           training=tf.convert_to_tensor(is_training))
            seqs_masks = tf.expand_dims(tf.cast(tf.not_equal(seqs_id, 0), tf.float32), -1)

        # multi-head attention
        seqs_outs = seqs_units * seqs_masks
        for i, (attention, dense) in enumerate(zip(self.list_attention, self.list_dense)):
            with tf.variable_scope("num_blocks_%d" % i):
                with tf.variable_scope("attention"):
                    seqs_outs = attention(layernorm(seqs_outs), seqs_outs, is_training, causality=True)

                with tf.variable_scope("feedforward"):
                    seqs_outs = dense(layernorm(seqs_outs), is_training)
                    seqs_outs *= seqs_masks

        with tf.variable_scope("output_ln"):
            seqs_outs = layernorm(seqs_outs)

        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def _sample_box_classifier_minibatch(self,
                                       proposal_boxlist,
                                       groundtruth_boxlist,
                                       groundtruth_classes_with_background):
    """Samples a mini-batch of proposals to be sent to the box classifier.

    Helper function for self._postprocess_rpn.

    Args:
      proposal_boxlist: A BoxList containing K proposal boxes in absolute
        coordinates.
      groundtruth_boxlist: A Boxlist containing N groundtruth object boxes in
        absolute coordinates.
      groundtruth_classes_with_background: A tensor with shape
        `[N, self.num_classes + 1]` representing groundtruth classes. The
        classes are assumed to be k-hot encoded, and include background as the
        zero-th class.

    Returns:
      a BoxList contained sampled proposals.
    """
    (cls_targets, cls_weights, _, _, _) = self._detector_target_assigner.assign(
        proposal_boxlist, groundtruth_boxlist,
        groundtruth_classes_with_background)
    # Selects all boxes as candidates if none of them is selected according
    # to cls_weights. This could happen as boxes within certain IOU ranges
    # are ignored. If triggered, the selected boxes will still be ignored
    # during loss computation.
    cls_weights += tf.to_float(tf.equal(tf.reduce_sum(cls_weights), 0))
    positive_indicator = tf.greater(tf.argmax(cls_targets, axis=1), 0)
    sampled_indices = self._second_stage_sampler.subsample(
        tf.cast(cls_weights, tf.bool),
        self._second_stage_batch_size,
        positive_indicator)
    return box_list_ops.boolean_mask(proposal_boxlist, sampled_indices)
//  [python]
def _sample_box_classifier_minibatch(self,
                                       proposal_boxlist,
                                       groundtruth_boxlist,
                                       groundtruth_classes_with_background):
    """Samples a mini-batch of proposals to be sent to the box classifier.

    Helper function for self._postprocess_rpn.

    Args:
      proposal_boxlist: A BoxList containing K proposal boxes in absolute
        coordinates.
      groundtruth_boxlist: A Boxlist containing N groundtruth object boxes in
        absolute coordinates.
      groundtruth_classes_with_background: A tensor with shape
        `[N, self.num_classes + 1]` representing groundtruth classes. The
        classes are assumed to be k-hot encoded, and include background as the
        zero-th class.

    Returns:
      a BoxList contained sampled proposals.
    """
    (cls_targets, cls_weights, _, _, _) = self._detector_target_assigner.assign(
        proposal_boxlist, groundtruth_boxlist,
        groundtruth_classes_with_background)
    # Selects all boxes as candidates if none of them is selected according
    # to cls_weights. This could happen as boxes within certain IOU ranges
    # are ignored. If triggered, the selected boxes will still be ignored
    # during loss computation.
    cls_weights += tf.cast(tf.equal(tf.reduce_sum(cls_weights), 0), tf.float32)
    positive_indicator = tf.greater(tf.argmax(cls_targets, axis=1), 0)
    sampled_indices = self._second_stage_sampler.subsample(
        tf.cast(cls_weights, tf.bool),
        self._second_stage_batch_size,
        positive_indicator)
    return box_list_ops.boolean_mask(proposal_boxlist, sampled_indices)
//  [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 _format_groundtruth_data(self, image_shape):
    """Helper function for preparing groundtruth data for target assignment.

    In order to be consistent with the model.DetectionModel interface,
    groundtruth boxes are specified in normalized coordinates and classes are
    specified as label indices with no assumed background category.  To prepare
    for target assignment, we:
    1) convert boxes to absolute coordinates,
    2) add a background class at class index 0
    3) groundtruth instance masks, if available, are resized to match
       image_shape.

    Args:
      image_shape: A 1-D int32 tensor of shape [4] representing the shape of the
        input image batch.

    Returns:
      groundtruth_boxlists: A list of BoxLists containing (absolute) coordinates
        of the groundtruth boxes.
      groundtruth_classes_with_background_list: A list of 2-D one-hot
        (or k-hot) tensors of shape [num_boxes, num_classes+1] containing the
        class targets with the 0th index assumed to map to the background class.
      groundtruth_masks_list: If present, a list of 3-D tf.float32 tensors of
        shape [num_boxes, image_height, image_width] containing instance masks.
        This is set to None if no masks exist in the provided groundtruth.
    """
    groundtruth_boxlists = [
        box_list_ops.to_absolute_coordinates(
            box_list.BoxList(boxes), image_shape[1], image_shape[2])
        for boxes in self.groundtruth_lists(fields.BoxListFields.boxes)]
    groundtruth_classes_with_background_list = [
        tf.to_float(
            tf.pad(one_hot_encoding, [[0, 0], [1, 0]], mode='CONSTANT'))
        for one_hot_encoding in self.groundtruth_lists(
            fields.BoxListFields.classes)]

    groundtruth_masks_list = self._groundtruth_lists.get(
        fields.BoxListFields.masks)
    if groundtruth_masks_list is not None:
      resized_masks_list = []
      for mask in groundtruth_masks_list:
        resized_4d_mask = tf.image.resize_images(
            tf.expand_dims(mask, axis=3),
            image_shape[1:3],
            method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
            align_corners=True)
        resized_masks_list.append(tf.squeeze(resized_4d_mask, axis=3))
      groundtruth_masks_list = resized_masks_list

    return (groundtruth_boxlists, groundtruth_classes_with_background_list,
            groundtruth_masks_list)
//  [python]
def _format_groundtruth_data(self, image_shape):
    """Helper function for preparing groundtruth data for target assignment.

    In order to be consistent with the model.DetectionModel interface,
    groundtruth boxes are specified in normalized coordinates and classes are
    specified as label indices with no assumed background category.  To prepare
    for target assignment, we:
    1) convert boxes to absolute coordinates,
    2) add a background class at class index 0
    3) groundtruth instance masks, if available, are resized to match
       image_shape.

    Args:
      image_shape: A 1-D int32 tensor of shape [4] representing the shape of the
        input image batch.

    Returns:
      groundtruth_boxlists: A list of BoxLists containing (absolute) coordinates
        of the groundtruth boxes.
      groundtruth_classes_with_background_list: A list of 2-D one-hot
        (or k-hot) tensors of shape [num_boxes, num_classes+1] containing the
        class targets with the 0th index assumed to map to the background class.
      groundtruth_masks_list: If present, a list of 3-D tf.float32 tensors of
        shape [num_boxes, image_height, image_width] containing instance masks.
        This is set to None if no masks exist in the provided groundtruth.
    """
    groundtruth_boxlists = [
        box_list_ops.to_absolute_coordinates(
            box_list.BoxList(boxes), image_shape[1], image_shape[2])
        for boxes in self.groundtruth_lists(fields.BoxListFields.boxes)]
    groundtruth_classes_with_background_list = [
        tf.cast(
            tf.pad(one_hot_encoding, [[0, 0], [1, 0]], mode='CONSTANT'), tf.float32)
        for one_hot_encoding in self.groundtruth_lists(
            fields.BoxListFields.classes)]

    groundtruth_masks_list = self._groundtruth_lists.get(
        fields.BoxListFields.masks)
    if groundtruth_masks_list is not None:
      resized_masks_list = []
      for mask in groundtruth_masks_list:
        resized_4d_mask = tf.image.resize_images(
            tf.expand_dims(mask, axis=3),
            image_shape[1:3],
            method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
            align_corners=True)
        resized_masks_list.append(tf.squeeze(resized_4d_mask, axis=3))
      groundtruth_masks_list = resized_masks_list

    return (groundtruth_boxlists, groundtruth_classes_with_background_list,
            groundtruth_masks_list)
//  [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 _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 `ugmentation_func applied to it.
    """
    image_height = tf.to_float(tf.shape(image)[0])
    image_width = tf.to_float(tf.shape(image)[1])
    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 _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 `ugmentation_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 noise_from_step_num():
  """Quantization noise equal to (phi * (step_num + 1)) mod 1.0.

  Not using random_uniform here due to a problem on TPU in that random seeds
  are not respected, which may cause the parameters on different replicas
  to go out-of-sync.

  Returns:
    a float32 scalar
  """
  step = tf.to_int32(tf.train.get_or_create_global_step()) + 1
  phi = ((5 ** 0.5) - 1) / 2
  # Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous
  # due to loss of precision when the step number gets large.
  # Computation in doubles does not work on TPU, so we use this complicated
  # alternative computation which does not suffer from these roundoff errors.
  ret = 0.0
  for i in range(30):
    ret += (((phi * (2 ** i)) % 1.0)  # double-precision computation in python
            * tf.to_float(tf.mod(step // (2 ** i), 2)))
  return tf.mod(ret, 1.0)
//  [python]
def noise_from_step_num():
  """Quantization noise equal to (phi * (step_num + 1)) mod 1.0.

  Not using random_uniform here due to a problem on TPU in that random seeds
  are not respected, which may cause the parameters on different replicas
  to go out-of-sync.

  Returns:
    a float32 scalar
  """
  step = tf.cast(tf.train.get_or_create_global_step(), tf.int32) + 1
  phi = ((5 ** 0.5) - 1) / 2
  # Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous
  # due to loss of precision when the step number gets large.
  # Computation in doubles does not work on TPU, so we use this complicated
  # alternative computation which does not suffer from these roundoff errors.
  ret = 0.0
  for i in range(30):
    ret += (((phi * (2 ** i)) % 1.0)  # double-precision computation in python
            * tf.to_float(tf.mod(step // (2 ** i), 2)))
  return tf.mod(ret, 1.0)
//  [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 _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 _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 add_keypoint_map(data):
    with tf.name_scope('add_keypoint_map'):
        image_shape = tf.shape(data['image'])[:2]
        kp = tf.minimum(tf.to_int32(tf.round(data['keypoints'])), image_shape-1)
        kmap = tf.scatter_nd(
                kp, tf.ones([tf.shape(kp)[0]], dtype=tf.int32), image_shape)
    return {**data, 'keypoint_map': kmap}
//  [python]
def add_keypoint_map(data):
    with tf.name_scope('add_keypoint_map'):
        image_shape = tf.shape(data['image'])[:2]
        kp = tf.minimum(tf.cast(tf.round(data['keypoints']), tf.int32), image_shape-1)
        kmap = tf.scatter_nd(
                kp, tf.ones([tf.shape(kp)[0]], dtype=tf.int32), image_shape)
    return {**data, 'keypoint_map': kmap}
//  [python]
def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
  """Round-off x to cand1 or to cand2 in an unbiased way.

  Cand1 and cand2 are the same shape as x.
  For every element of x, the corresponding elements of cand1 and cand2 should
  be the two closest bfloat16 values to x.  Order does not matter.
  cand1 and cand2 must differ from each other.

  Args:
    x: A float32 Tensor.
    noise: A Tensor broadcastable to the shape of x containing
    random uniform values in [0.0, 1.0].
    cand1: A bfloat16 Tensor the same shape as x.
    cand2: A bfloat16 Tensor the same shape as x.

  Returns:
    A bfloat16 Tensor.
  """
  cand1_f = tf.to_float(cand1)
  cand2_f = tf.to_float(cand2)
  step_size = cand2_f - cand1_f
  fpart = (x - cand1_f) / step_size
  ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
  return ret
//  [python]
def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
  """Round-off x to cand1 or to cand2 in an unbiased way.

  Cand1 and cand2 are the same shape as x.
  For every element of x, the corresponding elements of cand1 and cand2 should
  be the two closest bfloat16 values to x.  Order does not matter.
  cand1 and cand2 must differ from each other.

  Args:
    x: A float32 Tensor.
    noise: A Tensor broadcastable to the shape of x containing
    random uniform values in [0.0, 1.0].
    cand1: A bfloat16 Tensor the same shape as x.
    cand2: A bfloat16 Tensor the same shape as x.

  Returns:
    A bfloat16 Tensor.
  """
  cand1_f = tf.cast(cand1, tf.float32)
  cand2_f = tf.cast(cand2, tf.float32)
  step_size = cand2_f - cand1_f
  fpart = (x - cand1_f) / step_size
  ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
  return ret
//  [python]
def filter_groundtruth_with_nan_box_coordinates(tensor_dict):
  """Filters out groundtruth with no bounding boxes.

  Args:
    tensor_dict: a dictionary of following groundtruth tensors -
      fields.InputDataFields.groundtruth_boxes
      fields.InputDataFields.groundtruth_instance_masks
      fields.InputDataFields.groundtruth_classes
      fields.InputDataFields.groundtruth_is_crowd
      fields.InputDataFields.groundtruth_area
      fields.InputDataFields.groundtruth_label_types

  Returns:
    a dictionary of tensors containing only the groundtruth that have bounding
    boxes.
  """
  groundtruth_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
  nan_indicator_vector = tf.greater(tf.reduce_sum(tf.to_int32(
      tf.is_nan(groundtruth_boxes)), reduction_indices=[1]), 0)
  valid_indicator_vector = tf.logical_not(nan_indicator_vector)
  valid_indices = tf.where(valid_indicator_vector)

  return retain_groundtruth(tensor_dict, valid_indices)
//  [python]
def filter_groundtruth_with_nan_box_coordinates(tensor_dict):
  """Filters out groundtruth with no bounding boxes.

  Args:
    tensor_dict: a dictionary of following groundtruth tensors -
      fields.InputDataFields.groundtruth_boxes
      fields.InputDataFields.groundtruth_instance_masks
      fields.InputDataFields.groundtruth_classes
      fields.InputDataFields.groundtruth_is_crowd
      fields.InputDataFields.groundtruth_area
      fields.InputDataFields.groundtruth_label_types

  Returns:
    a dictionary of tensors containing only the groundtruth that have bounding
    boxes.
  """
  groundtruth_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
  nan_indicator_vector = tf.greater(tf.reduce_sum(tf.cast(
      tf.is_nan(groundtruth_boxes), tf.int32), reduction_indices=[1]), 0)
  valid_indicator_vector = tf.logical_not(nan_indicator_vector)
  valid_indices = tf.where(valid_indicator_vector)

  return retain_groundtruth(tensor_dict, valid_indices)
//  [python]
def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.

  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.to_bfloat16(x)
  cand1_f = tf.to_float(cand1)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.to_bfloat16(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.to_bfloat16(x_sign)
//  [python]
def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.

  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.cast(x, tf.bfloat16)
  cand1_f = tf.to_float(cand1)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.cast(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995), tf.bfloat16)
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.cast(x_sign, tf.bfloat16)
//  [python]
def ratio_preserving_resize(image, **config):
    target_size = tf.convert_to_tensor(config['resize'])
    scales = tf.to_float(tf.divide(target_size, tf.shape(image)[:2]))
    new_size = tf.to_float(tf.shape(image)[:2]) * tf.reduce_max(scales)
    image = tf.image.resize_images(image, tf.to_int32(new_size),
                                   method=tf.image.ResizeMethod.BILINEAR)
    return tf.image.resize_image_with_crop_or_pad(image, target_size[0], target_size[1])
//  [python]
def ratio_preserving_resize(image, **config):
    target_size = tf.convert_to_tensor(config['resize'])
    scales = tf.cast(tf.divide(target_size, tf.shape(image)[:2]), tf.float32)
    new_size = tf.cast(tf.shape(image)[:2], tf.float32) * tf.reduce_max(scales)
    image = tf.image.resize_images(image, tf.to_int32(new_size),
                                   method=tf.image.ResizeMethod.BILINEAR)
    return tf.image.resize_image_with_crop_or_pad(image, target_size[0], target_size[1])