jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def normalize_image(image, original_minval, original_maxval, target_minval,
                    target_maxval):
  """Normalizes pixel values in the image.

  Moves the pixel values from the current [original_minval, original_maxval]
  range to a the [target_minval, target_maxval] range.
  Args:
    image: rank 3 float32 tensor containing 1 image -> [height, width,
      channels].
    original_minval: current image minimum value.
    original_maxval: current image maximum value.
    target_minval: target image minimum value.
    target_maxval: target image maximum value.

  Returns:
    image: image which is the same shape as input image.
  """
  with tf.compat.v1.name_scope('NormalizeImage', values=[image]):
    original_minval = float(original_minval)
    original_maxval = float(original_maxval)
    target_minval = float(target_minval)
    target_maxval = float(target_maxval)
    image = tf.to_float(image)
    image = tf.subtract(image, original_minval)
    image = tf.multiply(image, (target_maxval - target_minval) /
                        (original_maxval - original_minval))
    image = tf.add(image, target_minval)
    return image
//  [python]
def normalize_image(image, original_minval, original_maxval, target_minval,
                    target_maxval):
  """Normalizes pixel values in the image.

  Moves the pixel values from the current [original_minval, original_maxval]
  range to a the [target_minval, target_maxval] range.
  Args:
    image: rank 3 float32 tensor containing 1 image -> [height, width,
      channels].
    original_minval: current image minimum value.
    original_maxval: current image maximum value.
    target_minval: target image minimum value.
    target_maxval: target image maximum value.

  Returns:
    image: image which is the same shape as input image.
  """
  with tf.compat.v1.name_scope('NormalizeImage', values=[image]):
    original_minval = float(original_minval)
    original_maxval = float(original_maxval)
    target_minval = float(target_minval)
    target_maxval = float(target_maxval)
    image = tf.cast(image, dtype=tf.float32)
    image = tf.subtract(image, original_minval)
    image = tf.multiply(image, (target_maxval - target_minval) /
                        (original_maxval - original_minval))
    image = tf.add(image, target_minval)
    return image
//  [python]
def _build_tree_from_leaf(leaf_nodes: tf.Tensor, arity: int) -> tf.RaggedTensor:
  """A function constructs a complete tree given all the leaf nodes.

  The function takes a 1-D array representing the leaf nodes of a tree and the
  tree's arity, and constructs a complete tree by recursively summing the
  adjacent children to get the parent until reaching the root node. Because we
  assume a complete tree, if the number of leaf nodes does not divide arity, the
  leaf nodes will be padded with zeros.

  Args:
    leaf_nodes: A 1-D array storing the leaf nodes of the tree.
    arity: A `int` for the branching factor of the tree, i.e. the number of
      children for each internal node.

  Returns:
    `tf.RaggedTensor` representing the tree. For example, if
    `leaf_nodes=tf.Tensor([1, 2, 3, 4])` and `arity=2`, then the returned value
    should be `tree=tf.RaggedTensor([[10],[3,7],[1,2,3,4]])`. In this way,
    `tree[layer][index]` can be used to access the node indexed by (layer,
    index) in the tree,
  """

  def pad_zero(leaf_nodes, size):
    paddings = tf.zeros(
        shape=(size - leaf_nodes.shape[0],), dtype=leaf_nodes.dtype)
    return tf.concat((leaf_nodes, paddings), axis=0)

  leaf_nodes_size = tf.constant(leaf_nodes.shape[0], dtype=tf.float32)
  num_layers = tf.math.ceil(
      tf.math.log(leaf_nodes_size) /
      tf.math.log(tf.cast(arity, dtype=tf.float32))) + 1
  leaf_nodes = pad_zero(
      leaf_nodes, tf.math.pow(tf.cast(arity, dtype=tf.float32), num_layers - 1))

  def _shrink_layer(layer: tf.Tensor, arity: int) -> tf.Tensor:
    return tf.reduce_sum((tf.reshape(layer, (-1, arity))), 1)

  # The following `tf.while_loop` constructs the tree from bottom up by
  # iteratively applying `_shrink_layer` to each layer of the tree. The reason
  # for the choice of TF1.0-style `tf.while_loop` is that @tf.function does not
  # support auto-translation from python loop to tf loop when loop variables
  # contain a `RaggedTensor` whose shape changes across iterations.

  idx = tf.identity(num_layers)
  loop_cond = lambda i, h: tf.less_equal(2.0, i)

  def _loop_body(i, h):
    return [
        tf.add(i, -1.0),
        tf.concat(([_shrink_layer(h[0], arity)], h), axis=0)
    ]

  _, tree = tf.while_loop(
      loop_cond,
      _loop_body, [idx, tf.RaggedTensor.from_tensor([leaf_nodes])],
      shape_invariants=[
          idx.get_shape(),
          tf.RaggedTensorSpec(dtype=leaf_nodes.dtype, ragged_rank=1)
      ])

  return tree
//  [python]
def get_grid_one_hot(box_gridy0y1, box_gridx0x1, feature_height, feature_width):
  """Get grid_one_hot from indices and feature_size."""
  (batch_size, num_boxes, output_size, _) = box_gridx0x1.get_shape().as_list()
  y_indices = tf.to_int32(
      tf.reshape(box_gridy0y1, [batch_size, num_boxes, output_size, 2]))
  x_indices = tf.to_int32(
      tf.reshape(box_gridx0x1, [batch_size, num_boxes, output_size, 2]))

  # shape is [batch_size, num_boxes, output_size, 2, height]
  grid_y_one_hot = tf.one_hot(tf.to_int32(y_indices), feature_height)
  # shape is [batch_size, num_boxes, output_size, 2, width]
  grid_x_one_hot = tf.one_hot(tf.to_int32(x_indices), feature_width)

  return grid_y_one_hot, grid_x_one_hot
//  [python]
def get_grid_one_hot(box_gridy0y1, box_gridx0x1, feature_height, feature_width):
  """Get grid_one_hot from indices and feature_size."""
  (batch_size, num_boxes, output_size, _) = box_gridx0x1.get_shape().as_list()
  y_indices = tf.cast(
      tf.reshape(box_gridy0y1, [batch_size, num_boxes, output_size, 2]),
      dtype=tf.int32)
  x_indices = tf.cast(
      tf.reshape(box_gridx0x1, [batch_size, num_boxes, output_size, 2]),
      dtype=tf.int32)

  # shape is [batch_size, num_boxes, output_size, 2, height]
  grid_y_one_hot = tf.one_hot(tf.cast(y_indices, tf.int32), feature_height)
  # shape is [batch_size, num_boxes, output_size, 2, width]
  grid_x_one_hot = tf.one_hot(tf.cast(x_indices, tf.int32), feature_width)

  return grid_y_one_hot, grid_x_one_hot
//  [python]
def __init__(self, shape, name=None):
        """Takes input in uint8 format which is cast to float32 and divided by 255
        before passing it to the model.

        On GPU this ensures lower data transfer times.

        Parameters
        ----------
        shape: [int]
            shape of the tensor.
        name: str
            name of the underlying placeholder
        """

        super().__init__(tf.placeholder(tf.uint8, [None] + list(shape), name=name))
        self._shape = shape
        self._output = tf.cast(super().get(), tf.float32) / 255.0
//  [python]
def create_mock_graph():
  g = tf.Graph()
  with g.as_default():
    in_image_tensor = tf.placeholder(
        tf.uint8, shape=[1, None, None, 3], name='image_tensor')
    tf.constant([2.0], name='num_detections')
    tf.constant(
        [[[0, 0.8, 0.7, 1], [0.1, 0.2, 0.8, 0.9], [0.2, 0.3, 0.4, 0.5]]],
        name='detection_boxes')
    tf.constant([[0.1, 0.2, 0.3]], name='detection_scores')
    tf.identity(
        tf.constant([[1.0, 2.0, 3.0]]) *
        tf.reduce_sum(tf.to_float(in_image_tensor)),
        name='detection_classes')
    graph_def = g.as_graph_def()

  with tf.gfile.Open(get_mock_graph_path(), 'w') as fl:
    fl.write(graph_def.SerializeToString())
//  [python]
def create_mock_graph():
  g = tf.Graph()
  with g.as_default():
    in_image_tensor = tf.placeholder(
        tf.uint8, shape=[1, None, None, 3], name='image_tensor')
    tf.constant([2.0], name='num_detections')
    tf.constant(
        [[[0, 0.8, 0.7, 1], [0.1, 0.2, 0.8, 0.9], [0.2, 0.3, 0.4, 0.5]]],
        name='detection_boxes')
    tf.constant([[0.1, 0.2, 0.3]], name='detection_scores')
    tf.identity(
        tf.constant([[1.0, 2.0, 3.0]]) *
        tf.reduce_sum(tf.cast(in_image_tensor, dtype=tf.float32)),
        name='detection_classes')
    graph_def = g.as_graph_def()

  with tf.gfile.Open(get_mock_graph_path(), 'w') as fl:
    fl.write(graph_def.SerializeToString())
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        sinusoidal_pos: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        if sinusoidal_pos is not None:
            if self.rotary_value:
                query_layer, key_layer, value_layer = self.apply_rotary_position_embeddings(
                    sinusoidal_pos, query_layer, key_layer, value_layer
                )
            else:
                query_layer, key_layer = self.apply_rotary_position_embeddings(sinusoidal_pos, query_layer, key_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.to_float(self.sqrt_att_head_size)
        attention_scores = tf.divide(attention_scores, dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFRoFormerModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        sinusoidal_pos: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        if sinusoidal_pos is not None:
            if self.rotary_value:
                query_layer, key_layer, value_layer = self.apply_rotary_position_embeddings(
                    sinusoidal_pos, query_layer, key_layer, value_layer
                )
            else:
                query_layer, key_layer = self.apply_rotary_position_embeddings(sinusoidal_pos, query_layer, key_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFRoFormerModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def _build(self, labels):
    # Cross-entropy.
    self._cross_entropy = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=labels, logits=self._predictor.logits))
    # Accuracy.
    correct_examples = tf.equal(labels, tf.argmax(self._predictor.logits, 1))
    self._nominal_accuracy = tf.reduce_mean(
        t.cast(correct_examples, tf.float32))

    # Interval bounds.
    if self._specification:
      bounds = self._specification(self._predictor.modules)
      v = tf.reduce_max(bounds, axis=1)
      self._interval_bounds_accuracy = tf.reduce_mean(
          t.cast(v <= 0., tf.float32))
      # Select specifications.
      if self._interval_bounds_loss_mode == 'all':
        pass  # Keep bounds the way it is.
      elif self._interval_bounds_loss_mode == 'most':
        bounds = tf.reduce_max(bounds, axis=1, keepdims=True)
      elif self._interval_bounds_loss_mode == 'random':
        idx = tf.random.uniform(
            [tf.shape(bounds)[0], self._interval_bounds_loss_n],
            0, tf.shape(bounds)[1], dtype=tf.int32)
        bounds = tf.batch_gather(bounds, idx)
      else:
        assert self._interval_bounds_loss_mode == 'least'
        # This picks the least violated contraint.
        mask = t.cast(bounds < 0., tf.float32)
        smallest_violation = tf.reduce_min(
            bounds + mask * _BIG_NUMBER, axis=1, keepdims=True)
        has_violations = tf.less(
            tf.reduce_sum(mask, axis=1, keepdims=True) + .5,
            t.cast(tf.shape(bounds)[1], tf.float32))
        largest_bounds = tf.reduce_max(bounds, axis=1, keepdims=True)
        bounds = tf.where(has_violations, smallest_violation, largest_bounds)

      if self._interval_bounds_loss_type == 'xent':
        v = tf.concat(
            [bounds, tf.zeros([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],
            axis=1)
        l = tf.concat(
            [tf.zeros_like(bounds),
             tf.ones([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],
            axis=1)
        self._verified_loss = tf.reduce_mean(
            tf.nn.softmax_cross_entropy_with_logits_v2(
                labels=tf.stop_gradient(l), logits=v))
      elif self._interval_bounds_loss_type == 'softplus':
        self._verified_loss = tf.reduce_mean(
            tf.nn.softplus(bounds + self._interval_bounds_hinge_margin))
      else:
        assert self._interval_bounds_loss_type == 'hinge'
        self._verified_loss = tf.reduce_mean(
            tf.maximum(bounds, -self._interval_bounds_hinge_margin))
    else:
      self._verified_loss = tf.constant(0.)
      self._interval_bounds_accuracy = tf.constant(0.)

    # PGD attack.
    if self._attack:
      self._attack(labels)
      self._attack_accuracy = self._attack.accuracy
      self._attack_cross_entropy = tf.reduce_mean(
          tf.nn.sparse_softmax_cross_entropy_with_logits(
              labels=labels, logits=self._attack.logits))
    else:
      self._attack_accuracy = tf.constant(0.)
      self._attack_cross_entropy = tf.constant(0.)
//  [python]
def _build(self, labels):
    # Cross-entropy.
    self._cross_entropy = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=labels, logits=self._predictor.logits))
    # Accuracy.
    correct_examples = tf.equal(labels, tf.argmax(self._predictor.logits, 1))
    self._nominal_accuracy = tf.reduce_mean(
        tf.cast(correct_examples, tf.float32))

    # Interval bounds.
    if self._specification:
      bounds = self._specification(self._predictor.modules)
      v = tf.reduce_max(bounds, axis=1)
      self._interval_bounds_accuracy = tf.reduce_mean(
          tf.cast(v <= 0., tf.float32))
      # Select specifications.
      if self._interval_bounds_loss_mode == 'all':
        pass  # Keep bounds the way it is.
      elif self._interval_bounds_loss_mode == 'most':
        bounds = tf.reduce_max(bounds, axis=1, keepdims=True)
      elif self._interval_bounds_loss_mode == 'random':
        idx = tf.random.uniform(
            [tf.shape(bounds)[0], self._interval_bounds_loss_n],
            0, tf.shape(bounds)[1], dtype=tf.int32)
        bounds = tf.batch_gather(bounds, idx)
      else:
        assert self._interval_bounds_loss_mode == 'least'
        # This picks the least violated contraint.
        mask = tf.cast(bounds < 0., tf.float32)
        smallest_violation = tf.reduce_min(
            bounds + mask * _BIG_NUMBER, axis=1, keepdims=True)
        has_violations = tf.less(
            tf.reduce_sum(mask, axis=1, keepdims=True) + .5,
            tf.cast(tf.shape(bounds)[1], tf.float32))
        largest_bounds = tf.reduce_max(bounds, axis=1, keepdims=True)
        bounds = tf.where(has_violations, smallest_violation, largest_bounds)

      if self._interval_bounds_loss_type == 'xent':
        v = tf.concat(
            [bounds, tf.zeros([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],
            axis=1)
        l = tf.concat(
            [tf.zeros_like(bounds),
             tf.ones([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],
            axis=1)
        self._verified_loss = tf.reduce_mean(
            tf.nn.softmax_cross_entropy_with_logits_v2(
                labels=tf.stop_gradient(l), logits=v))
      elif self._interval_bounds_loss_type == 'softplus':
        self._verified_loss = tf.reduce_mean(
            tf.nn.softplus(bounds + self._interval_bounds_hinge_margin))
      else:
        assert self._interval_bounds_loss_type == 'hinge'
        self._verified_loss = tf.reduce_mean(
            tf.maximum(bounds, -self._interval_bounds_hinge_margin))
    else:
      self._verified_loss = tf.constant(0.)
      self._interval_bounds_accuracy = tf.constant(0.)

    # PGD attack.
    if self._attack:
      self._attack(labels)
      self._attack_accuracy = self._attack.accuracy
      self._attack_cross_entropy = tf.reduce_mean(
          tf.nn.sparse_softmax_cross_entropy_with_logits(
              labels=labels, logits=self._attack.logits))
    else:
      self._attack_accuracy = tf.constant(0.)
      self._attack_cross_entropy = tf.constant(0.)
//  [python]
def mrcnn_class_loss_graph(target_class_ids, pred_class_logits,
                           active_class_ids):
    """Loss for the classifier head of Mask RCNN.

    target_class_ids: [batch, num_rois]. Integer class IDs. Uses zero
        padding to fill in the array.
    pred_class_logits: [batch, num_rois, num_classes]
    active_class_ids: [batch, num_classes]. Has a value of 1 for
        classes that are in the dataset of the image, and 0
        for classes that are not in the dataset.
    """
    # During model building, Keras calls this function with
    # target_class_ids of type float32. Unclear why. Cast it
    # to int to get around it.
    target_class_ids = tf.to_int64(target_class_ids)

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(pred_class_logits, axis=2)
    # TODO: Update this line to work with batch > 1. Right now it assumes all
    #       images in a batch have the same active_class_ids
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    loss = tf.reduce_sum(loss) / tf.reduce_sum(pred_active)
    return loss
//  [python]
def mrcnn_class_loss_graph(target_class_ids, pred_class_logits,
                           active_class_ids):
    """Loss for the classifier head of Mask RCNN.

    target_class_ids: [batch, num_rois]. Integer class IDs. Uses zero
        padding to fill in the array.
    pred_class_logits: [batch, num_rois, num_classes]
    active_class_ids: [batch, num_classes]. Has a value of 1 for
        classes that are in the dataset of the image, and 0
        for classes that are not in the dataset.
    """
    # During model building, Keras calls this function with
    # target_class_ids of type float32. Unclear why. Cast it
    # to int to get around it.
    target_class_ids = tf.cast(target_class_ids, 'int64')

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(pred_class_logits, axis=2)
    # TODO: Update this line to work with batch > 1. Right now it assumes all
    #       images in a batch have the same active_class_ids
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    loss = tf.reduce_sum(loss) / tf.reduce_sum(pred_active)
    return loss
//  [python]
def trim_zeros_graph(boxes, name=None):
    """Often boxes are represented with matricies of shape [N, 4] and
    are padded with zeros. This removes zero boxes.

    boxes: [N, 4] matrix of boxes.
    non_zeros: [N] a 1D boolean mask identifying the rows to keep
    """
    non_zeros = tf.cast(tf.reduce_sum(tf.abs(boxes), axis=1), tf.bool)
    boxes = tf.boolean_mask(boxes, non_zeros, name=name)
    return boxes, non_zeros
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indicies.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_roi_ix))
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(tf.size(target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indicies.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_roi_ix), tf.int64)
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(tf.size(target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def single_level_feature_crop(features, level_boxes, detection_prior_levels,
                              min_mask_level, mask_crop_size):
  """Crop the FPN features at the appropriate levels for each detection.


  Args:
    features: a float tensor of shape [batch_size, num_levels, max_feature_size,
      max_feature_size, num_downsample_channels].
    level_boxes: a float Tensor of the level boxes to crop from. [batch_size,
      num_instances, 4].
    detection_prior_levels: an int Tensor of instance assigned level of shape
      [batch_size, num_instances].
    min_mask_level: minimum FPN level to crop mask feature from.
    mask_crop_size: an int of mask crop size.

  Returns:
    crop_features: a float Tensor of shape [batch_size * num_instances,
        mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
        instance feature crop.
  """
  (batch_size, num_levels, max_feature_size, _,
   num_downsample_channels) = features.get_shape().as_list()
  _, num_of_instances, _ = level_boxes.get_shape().as_list()
  level_boxes = tf.to_int32(level_boxes)
  assert num_of_instances == detection_prior_levels.get_shape().as_list()[1]

  x_start_indices = level_boxes[:, :, 1]
  y_start_indices = level_boxes[:, :, 0]
  # generate the full indices (not just the starting index)
  x_idx_list = []
  y_idx_list = []
  for i in range(mask_crop_size):
    x_idx_list.append(x_start_indices + i)
    y_idx_list.append(y_start_indices + i)

  x_indices = tf.stack(x_idx_list, axis=2)
  y_indices = tf.stack(y_idx_list, axis=2)
  levels = detection_prior_levels - min_mask_level
  height_dim_size = max_feature_size
  level_dim_size = max_feature_size * height_dim_size
  batch_dim_size = num_levels * level_dim_size
  # TODO(weicheng) change this to gather_nd for better readability.
  indices = tf.reshape(
      tf.tile(
          tf.reshape(
              tf.range(batch_size) * batch_dim_size, [batch_size, 1, 1, 1]),
          [1, num_of_instances, mask_crop_size, mask_crop_size]) + tf.tile(
              tf.reshape(levels * level_dim_size,
                         [batch_size, num_of_instances, 1, 1]),
              [1, 1, mask_crop_size, mask_crop_size]) + tf.tile(
                  tf.reshape(y_indices * height_dim_size,
                             [batch_size, num_of_instances, mask_crop_size, 1]),
                  [1, 1, 1, mask_crop_size]) +
      tf.tile(
          tf.reshape(x_indices,
                     [batch_size, num_of_instances, 1, mask_crop_size]),
          [1, 1, mask_crop_size, 1]), [-1])

  features_r2 = tf.reshape(features, [-1, num_downsample_channels])
  crop_features = tf.reshape(
      tf.gather(features_r2, indices), [
          batch_size * num_of_instances, mask_crop_size, mask_crop_size,
          num_downsample_channels
      ])

  return crop_features
//  [python]
def single_level_feature_crop(features, level_boxes, detection_prior_levels,
                              min_mask_level, mask_crop_size):
  """Crop the FPN features at the appropriate levels for each detection.


  Args:
    features: a float tensor of shape [batch_size, num_levels, max_feature_size,
      max_feature_size, num_downsample_channels].
    level_boxes: a float Tensor of the level boxes to crop from. [batch_size,
      num_instances, 4].
    detection_prior_levels: an int Tensor of instance assigned level of shape
      [batch_size, num_instances].
    min_mask_level: minimum FPN level to crop mask feature from.
    mask_crop_size: an int of mask crop size.

  Returns:
    crop_features: a float Tensor of shape [batch_size * num_instances,
        mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
        instance feature crop.
  """
  (batch_size, num_levels, max_feature_size, _,
   num_downsample_channels) = features.get_shape().as_list()
  _, num_of_instances, _ = level_boxes.get_shape().as_list()
  level_boxes = tf.cast(level_boxes, tf.int32)
  assert num_of_instances == detection_prior_levels.get_shape().as_list()[1]

  x_start_indices = level_boxes[:, :, 1]
  y_start_indices = level_boxes[:, :, 0]
  # generate the full indices (not just the starting index)
  x_idx_list = []
  y_idx_list = []
  for i in range(mask_crop_size):
    x_idx_list.append(x_start_indices + i)
    y_idx_list.append(y_start_indices + i)

  x_indices = tf.stack(x_idx_list, axis=2)
  y_indices = tf.stack(y_idx_list, axis=2)
  levels = detection_prior_levels - min_mask_level
  height_dim_size = max_feature_size
  level_dim_size = max_feature_size * height_dim_size
  batch_dim_size = num_levels * level_dim_size
  # TODO(weicheng) change this to gather_nd for better readability.
  indices = tf.reshape(
      tf.tile(
          tf.reshape(
              tf.range(batch_size) * batch_dim_size, [batch_size, 1, 1, 1]),
          [1, num_of_instances, mask_crop_size, mask_crop_size]) + tf.tile(
              tf.reshape(levels * level_dim_size,
                         [batch_size, num_of_instances, 1, 1]),
              [1, 1, mask_crop_size, mask_crop_size]) + tf.tile(
                  tf.reshape(y_indices * height_dim_size,
                             [batch_size, num_of_instances, mask_crop_size, 1]),
                  [1, 1, 1, mask_crop_size]) +
      tf.tile(
          tf.reshape(x_indices,
                     [batch_size, num_of_instances, 1, mask_crop_size]),
          [1, 1, mask_crop_size, 1]), [-1])

  features_r2 = tf.reshape(features, [-1, num_downsample_channels])
  crop_features = tf.reshape(
      tf.gather(features_r2, indices), [
          batch_size * num_of_instances, mask_crop_size, mask_crop_size,
          num_downsample_channels
      ])

  return crop_features
//  [python]
def create_position_ids_from_input_ids(self, input_ids):
        """
        Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding
        symbols are ignored. This is modified from fairseq's `utils.make_positions`.

        Args:
            input_ids: tf.Tensor
        Returns: tf.Tensor
        """
        mask = tf.to_int32(tf.math.not_equal(input_ids, self.padding_idx))
        incremental_indices = tf.math.cumsum(mask, axis=1) * mask

        return incremental_indices + self.padding_idx
//  [python]
def create_position_ids_from_input_ids(self, input_ids):
        """
        Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding
        symbols are ignored. This is modified from fairseq's `utils.make_positions`.

        Args:
            input_ids: tf.Tensor
        Returns: tf.Tensor
        """
        mask = tf.cast(tf.math.not_equal(input_ids, self.padding_idx), dtype=input_ids.dtype)
        incremental_indices = tf.math.cumsum(mask, axis=1) * mask

        return incremental_indices + self.padding_idx
//  [python]
def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_ix))
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def denorm_boxes_graph(boxes, shape):
    """Converts boxes from normalized coordinates to pixel coordinates.
    boxes: [..., (y1, x1, y2, x2)] in normalized coordinates
    shape: [..., (height, width)] in pixels

    Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
    coordinates it's inside the box.

    Returns:
        [..., (y1, x1, y2, x2)] in pixel coordinates
    """
    h, w = tf.split(t.cast(shape, tf.float32), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return t.cast(tf.round(tf.multiply(boxes, scale) + shift), tf.int32)
//  [python]
def denorm_boxes_graph(boxes, shape):
    """Converts boxes from normalized coordinates to pixel coordinates.
    boxes: [..., (y1, x1, y2, x2)] in normalized coordinates
    shape: [..., (height, width)] in pixels

    Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
    coordinates it's inside the box.

    Returns:
        [..., (y1, x1, y2, x2)] in pixel coordinates
    """
    h, w = tf.split(tf.cast(shape, tf.float32), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return tf.cast(tf.round(tf.multiply(boxes, scale) + shift), tf.int32)
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=self.config.GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.cast(tf.size(w), tf.float32)
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def _test_forward_rel_op(data, func):
    with tf.Graph().as_default():
        in1 = tf.placeholder(shape=data[0].shape, dtype=data[0].dtype, name="in1")
        in2 = tf.placeholder(shape=data[1].shape, dtype=data[1].dtype, name="in2")
        op = func(in1, in2, name="op")
        out = tf.cast(op, tf.int32, name="out1")
        compare_tf_with_tvm([data[0], data[1]], ["in1:0", "in2:0"], "out1:0")
//  [python]
def _from_numpy(self, a, type_as=None):
        if isinstance(a, float):
            a = np.array(a)
        if not isinstance(a, self.__type__):
            if type_as is None:
                return tf.convert_to_tensor(a)
            else:
                return tf.convert_to_tensor(a, dtype=type_as.dtype)
        else:
            if type_as is None:
                return a
            else:
                return tf.cast(a, dtype=type_as.dtype)
//  [python]
def to_textline_dataset(files_with_sizes, config):

    def parse_csv(value):
        cont_defaults = [[0.0] for i in range(1, 14)]
        cate_defaults = [[0] for i in range(1, 27)]
        label_defaults = [[0]]
        column_headers = TRAIN_DATA_COLUMNS
        record_defaults = label_defaults + cont_defaults + cate_defaults
        columns = tf.io.decode_csv(value, record_defaults=record_defaults)
        all_columns = collections.OrderedDict(zip(column_headers, columns))
        labels = all_columns.pop(LABEL_COLUMN[0])
        features = all_columns
        # EmbeddingVariable only supports int64 not int32.
        for j in range(len(CATEGORICAL_COLUMNS)):
            features[CATEGORICAL_COLUMNS[j]] = \
                tf.to_int64(features[CATEGORICAL_COLUMNS[j]])
        return features, labels

    print(files_with_sizes)
    files = [pair[0] for pair in files_with_sizes]
    data_size = sum([pair[1] for pair in files_with_sizes])

    dataset = tf.data.TextLineDataset(files)
    dataset = dataset.shuffle(buffer_size=20000,
                              seed=config["seed"])  # fix seed for reproducing
    dataset = dataset.repeat(config["no_of_epochs"])
    dataset = dataset.prefetch(config["batch_size"])
    dataset = dataset.batch(config["batch_size"])
    dataset = dataset.map(parse_csv, num_parallel_calls=28)
    dataset = dataset.prefetch(2)
    return dataset, data_size
//  [python]
def to_textline_dataset(files_with_sizes, config):

    def parse_csv(value):
        cont_defaults = [[0.0] for i in range(1, 14)]
        cate_defaults = [[0] for i in range(1, 27)]
        label_defaults = [[0]]
        column_headers = TRAIN_DATA_COLUMNS
        record_defaults = label_defaults + cont_defaults + cate_defaults
        columns = tf.io.decode_csv(value, record_defaults=record_defaults)
        all_columns = collections.OrderedDict(zip(column_headers, columns))
        labels = all_columns.pop(LABEL_COLUMN[0])
        features = all_columns
        # EmbeddingVariable only supports int64 not int32.
        for j in range(len(CATEGORICAL_COLUMNS)):
            features[CATEGORICAL_COLUMNS[j]] = \
                tf.cast(features[CATEGORICAL_COLUMNS[j]], dtype=tf.int64)
        return features, labels

    print(files_with_sizes)
    files = [pair[0] for pair in files_with_sizes]
    data_size = sum([pair[1] for pair in files_with_sizes])

    dataset = tf.data.TextLineDataset(files)
    dataset = dataset.shuffle(buffer_size=20000,
                              seed=config["seed"])  # fix seed for reproducing
    dataset = dataset.repeat(config["no_of_epochs"])
    dataset = dataset.prefetch(config["batch_size"])
    dataset = dataset.batch(config["batch_size"])
    dataset = dataset.map(parse_csv, num_parallel_calls=28)
    dataset = dataset.prefetch(2)
    return dataset, data_size
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.to_float(self.sqrt_att_head_size)
        attention_scores = tf.divide(attention_scores, dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFRobertaModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFRobertaModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs