jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, x, training, mask):

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.sqrt(tf.to_float(self.d_model))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x
//  [python]
def call(self, x, training, mask):

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x
//  [python]
def preprocess_for_quantization(image_data, image_size, crop_padding=32):
  """Crops to center of image with padding then scales, normalizes image_size.

  Args:
    image_data: A 3D Tensor representing the RGB image data. Image can be of
      arbitrary height and width.
    image_size: image height/width dimension.
    crop_padding: the padding size to use when centering the crop.

  Returns:
    A decoded and cropped image Tensor. Image is normalized to [-1,1].
  """
  shape = tf.shape(image_data)
  image_height = shape[0]
  image_width = shape[1]

  padded_center_crop_size = tf.to_int32(
      (image_size * 1.0 / (image_size + crop_padding)) *
      tf.to_float(tf.minimum(image_height, image_width)))

  offset_height = ((image_height - padded_center_crop_size) + 1) // 2
  offset_width = ((image_width - padded_center_crop_size) + 1) // 2

  image = tf.image.crop_to_bounding_box(
      image_data,
      offset_height=offset_height,
      offset_width=offset_width,
      target_height=padded_center_crop_size,
      target_width=padded_center_crop_size)

  image = tf.image.resize([image], [image_size, image_size],
                          method=tf.image.ResizeMethod.BILINEAR)[0]
  image = tf.to_float(image)
  image -= tf.constant(MEAN_RGB)
  image /= tf.constant(STDDEV_RGB)
  return image
//  [python]
def preprocess_for_quantization(image_data, image_size, crop_padding=32):
  """Crops to center of image with padding then scales, normalizes image_size.

  Args:
    image_data: A 3D Tensor representing the RGB image data. Image can be of
      arbitrary height and width.
    image_size: image height/width dimension.
    crop_padding: the padding size to use when centering the crop.

  Returns:
    A decoded and cropped image Tensor. Image is normalized to [-1,1].
  """
  shape = tf.shape(image_data)
  image_height = shape[0]
  image_width = shape[1]

  padded_center_crop_size = tf.cast(
      (image_size * 1.0 / (image_size + crop_padding)) *
      tf.cast(tf.minimum(image_height, image_width), tf.float32), tf.int32)

  offset_height = ((image_height - padded_center_crop_size) + 1) // 2
  offset_width = ((image_width - padded_center_crop_size) + 1) // 2

  image = tf.image.crop_to_bounding_box(
      image_data,
      offset_height=offset_height,
      offset_width=offset_width,
      target_height=padded_center_crop_size,
      target_width=padded_center_crop_size)

  image = tf.image.resize([image], [image_size, image_size],
                          method=tf.image.ResizeMethod.BILINEAR)[0]
  image = tf.cast(image, tf.float32)
  image -= tf.constant(MEAN_RGB)
  image /= tf.constant(STDDEV_RGB)
  return image
//  [python]
def _aggregate_text_embedding(self, token_ids, embeddings):
    """Aggregate text embedding for a UI element."""
    if self._hparams['obj_text_aggregation'] == 'max':
      # Find valid tokens (not PADDING/EOS/UNK/START).
      valid_token_mask = tf.greater_equal(token_ids, 4)
      # Use large negative bias for invalid tokens.
      invalid_token_bias = t.cast(
          tf.logical_not(valid_token_mask), t.float32) * -1e9
      # [batch, node_num, word_num, hidden_size]
      embeddings = embeddings + tf.expand_dims(invalid_token_bias, axis=-1)
      # Max value for each dimension, [batch, node_num, hidden_size].
      embeddings = tf.reduce_max(embeddings, axis=-2)
      # For objects with no text, use 0.
      valid_object_mask = t.cast(
          tf.reduce_any(valid_token_mask, axis=-1), t.float32)
      embeddings = embeddings * tf.expand_dims(valid_object_mask, axis=-1)

    elif self._hparams['obj_text_aggregation'] == 'sum':
      # [batch, step, #max_obj, #max_token]  0 for padded tokens
      real_objects = t.cast(tf.greater_equal(token_ids, 4), t.float32)
      # [batch, step, #max_obj, hidden]   0s for padded objects
      embeddings = tf.reduce_sum(
          input_tensor=embeddings * tf.expand_dims(real_objects, 3), axis=-2)

    else:
      raise ValueError('Unrecognized token aggregation %s' %
                       (self._hparams['obj_text_aggregation']))
    return embeddings
//  [python]
def _aggregate_text_embedding(self, token_ids, embeddings):
    """Aggregate text embedding for a UI element."""
    if self._hparams['obj_text_aggregation'] == 'max':
      # Find valid tokens (not PADDING/EOS/UNK/START).
      valid_token_mask = tf.greater_equal(token_ids, 4)
      # Use large negative bias for invalid tokens.
      invalid_token_bias = tf.cast(
          tf.logical_not(valid_token_mask), tf.float32) * -1e9
      # [batch, node_num, word_num, hidden_size]
      embeddings = embeddings + tf.expand_dims(invalid_token_bias, axis=-1)
      # Max value for each dimension, [batch, node_num, hidden_size].
      embeddings = tf.reduce_max(embeddings, axis=-2)
      # For objects with no text, use 0.
      valid_object_mask = tf.cast(
          tf.reduce_any(valid_token_mask, axis=-1), tf.float32)
      embeddings = embeddings * tf.expand_dims(valid_object_mask, axis=-1)

    elif self._hparams['obj_text_aggregation'] == 'sum':
      # [batch, step, #max_obj, #max_token]  0 for padded tokens
      real_objects = tf.cast(tf.greater_equal(token_ids, 4), tf.float32)
      # [batch, step, #max_obj, hidden]   0s for padded objects
      embeddings = tf.reduce_sum(
          input_tensor=embeddings * tf.expand_dims(real_objects, 3), axis=-2)

    else:
      raise ValueError('Unrecognized token aggregation %s' %
                       (self._hparams['obj_text_aggregation']))
    return embeddings
//  [python]
def _cross_replica_average(self, t: tf.Tensor, num_shards_per_group: int):
    """Calculates the average value of input tensor across TPU replicas."""
    num_shards = tpu_function.get_tpu_context().number_of_shards
    group_assignment = None
    if num_shards_per_group > 1:
      if num_shards % num_shards_per_group != 0:
        raise ValueError(
            'num_shards: %d mod shards_per_group: %d, should be 0' %
            (num_shards, num_shards_per_group))
      num_groups = num_shards // num_shards_per_group
      group_assignment = [[
          x for x in range(num_shards) if x // num_shards_per_group == y
      ] for y in range(num_groups)]
    return tf1.tpu.cross_replica_sum(t, group_assignment) / t.dtype.max
//  [python]
def _cross_replica_average(self, t: tf.Tensor, num_shards_per_group: int):
    """Calculates the average value of input tensor across TPU replicas."""
    num_shards = tpu_function.get_tpu_context().number_of_shards
    group_assignment = None
    if num_shards_per_group > 1:
      if num_shards % num_shards_per_group != 0:
        raise ValueError(
            'num_shards: %d mod shards_per_group: %d, should be 0' %
            (num_shards, num_shards_per_group))
      num_groups = num_shards // num_shards_per_group
      group_assignment = [[
          x for x in range(num_shards) if x // num_shards_per_group == y
      ] for y in range(num_groups)]
    return tf1.tpu.cross_replica_sum(t, group_assignment) / tf.cast(
        num_shards_per_group, t.dtype)
//  [python]
def get_data(self):
        sample = self.iterator.get_next()
        self.user = sample['user']
        self.item = sample['item']
        self.label = tf.to_float(sample['label'])
//  [python]
def get_data(self):
        sample = self.iterator.get_next()
        self.user = sample['user']
        self.item = sample['item']
        self.label = tf.cast(sample['label'],tf.float32)
//  [python]
def _prepare_local(self, var_device, var_dtype, apply_state):
    super()._prepare_local(var_device, var_dtype, apply_state)

    local_step = t.cast(self.iterations + 1, var_dtype)
    beta_1_t = tf.identity(self._get_hyper("beta_1", var_dtype))
    beta_2_t = tf.identity(self._get_hyper("beta_2", var_dtype))
    weight_decay_rate = tf.identity(
        self._get_hyper("weight_decay_rate", var_dtype)
    )
    beta_1_power = tf.pow(beta_1_t, local_step)
    beta_2_power = tf.pow(beta_2_t, local_step)
    apply_state[(var_device, var_dtype)].update(
        dict(
            weight_decay_rate=weight_decay_rate,
            epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),
            beta_1_t=beta_1_t,
            beta_1_power=beta_1_power,
            one_minus_beta_1_t=1 - beta_1_t,
            beta_2_t=beta_2_t,
            beta_2_power=beta_2_power,
            one_minus_beta_2_t=1 - beta_2_t,
        )
    )
//  [python]
def _prepare_local(self, var_device, var_dtype, apply_state):
    super()._prepare_local(var_device, var_dtype, apply_state)

    local_step = tf.cast(self.iterations + 1, var_dtype)
    beta_1_t = tf.identity(self._get_hyper("beta_1", var_dtype))
    beta_2_t = tf.identity(self._get_hyper("beta_2", var_dtype))
    weight_decay_rate = tf.identity(
        self._get_hyper("weight_decay_rate", var_dtype)
    )
    beta_1_power = tf.pow(beta_1_t, local_step)
    beta_2_power = tf.pow(beta_2_t, local_step)
    apply_state[(var_device, var_dtype)].update(
        dict(
            weight_decay_rate=weight_decay_rate,
            epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),
            beta_1_t=beta_1_t,
            beta_1_power=beta_1_power,
            one_minus_beta_1_t=1 - beta_1_t,
            beta_2_t=beta_2_t,
            beta_2_power=beta_2_power,
            one_minus_beta_2_t=1 - beta_2_t,
        )
    )
//  [python]
def build_losses(self,
                   labels,
                   model_outputs,
                   metrics,
                   aux_losses=None) -> Tuple[tf.Tensor, tf.Tensor]:
    """Return the final loss, and the masked-lm loss."""
    with tf.name_scope('MaskedLMTask/losses'):
      metrics = dict([(metric.name, metric) for metric in metrics])
      lm_prediction_losses = tf_keras.losses.sparse_categorical_crossentropy(
          labels['masked_lm_ids'],
          tf.cast(model_outputs['mlm_logits'], tf.float32),
          from_logits=True)
      lm_label_weights = labels['masked_lm_weights']
      lm_numerator_loss = tf.reduce_sum(lm_prediction_losses *
                                        lm_label_weights)
      lm_denominator_loss = tf.reduce_sum(lm_label_weights)
      mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
      metrics['lm_example_loss'].update_state(mlm_loss)
      if 'next_sentence_labels' in labels:
        sentence_labels = labels['next_sentence_labels']
        sentence_outputs = tf.cast(
            model_outputs['next_sentence'], dtype=tf.float32)
        sentence_loss = tf.reduce_mean(
            tf_keras.losses.sparse_categorical_crossentropy(
                sentence_labels, sentence_outputs, from_logits=True))
        metrics['next_sentence_loss'].update_state(sentence_loss)
        total_loss = mlm_loss + sentence_loss
      else:
        total_loss = mlm_loss

      if aux_losses:
        total_loss += tf.add_n(aux_losses)
      return total_loss, lm_prediction_losses
//  [python]
def build_fcn_net(self, inp, use_dice = False):
        bn1 = tf.layers.batch_normalization(inputs=inp, name='bn1')
        dnn1 = tf.layers.dense(bn1, 200, activation=None, name='f1')
        if use_dice:
            dnn1 = dice(dnn1, name='dice_1')
        else:
            dnn1 = prelu(dnn1, scope='prelu_1')

        dnn2 = tf.layers.dense(dnn1, 80, activation=None, name='f2')
        if use_dice:
            dnn2 = dice(dnn2, name='dice_2')
        else:
            dnn2 = prelu(dnn2, scope='prelu_2')

        dnn3 = tf.layers.dense(dnn2, 2, activation=None, name='f3')
        self.y_hat = tf.nn.softmax(dnn3) + 0.00000001

        with tf.name_scope('Metrics'):
            # Cross-entropy loss and optimizer initialization
            ctr_loss = - tf.reduce_mean(tf.log(self.y_hat) * self.target_ph)
            self.loss = ctr_loss
            if self.use_negsample:
                self.loss += self.aux_loss
            if self.reg:
                self.loss += self.reg_loss

            tf.summary.scalar('loss', self.loss)
            self.optimizer = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)
            # Accuracy metric
            self.accuracy = tf.reduce_mean(tf.to_float(tf.equal(tf.round(self.y_hat), self.target_ph)))
            tf.summary.scalar('accuracy', self.accuracy)

        self.merged = tf.summary.merge_all()
//  [python]
def build_fcn_net(self, inp, use_dice = False):
        bn1 = tf.layers.batch_normalization(inputs=inp, name='bn1')
        dnn1 = tf.layers.dense(bn1, 200, activation=None, name='f1')
        if use_dice:
            dnn1 = dice(dnn1, name='dice_1')
        else:
            dnn1 = prelu(dnn1, scope='prelu_1')

        dnn2 = tf.layers.dense(dnn1, 80, activation=None, name='f2')
        if use_dice:
            dnn2 = dice(dnn2, name='dice_2')
        else:
            dnn2 = prelu(dnn2, scope='prelu_2')

        dnn3 = tf.layers.dense(dnn2, 2, activation=None, name='f3')
        self.y_hat = tf.nn.softmax(dnn3) + 0.00000001

        with tf.name_scope('Metrics'):
            # Cross-entropy loss and optimizer initialization
            ctr_loss = - tf.reduce_mean(tf.log(self.y_hat) * self.target_ph)
            self.loss = ctr_loss
            if self.use_negsample:
                self.loss += self.aux_loss
            if self.reg:
                self.loss += self.reg_loss

            tf.summary.scalar('loss', self.loss)
            self.optimizer = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)
            # Accuracy metric
            self.accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.round(self.y_hat), self.target_ph), tf.float32))
            tf.summary.scalar('accuracy', self.accuracy)

        self.merged = tf.summary.merge_all()
//  [python]
def smooth_l1_loss_rcnn(bbox_targets, bbox_pred, anchor_state, sigma=3.0):

    outside_mask = tf.stop_gradient(tf.to_float(tf.greater(anchor_state, 0)))

    bbox_pred = tf.reshape(bbox_pred, [-1, 1, 4])
    bbox_targets = tf.reshape(bbox_targets, [-1, 1, 4])

    value = _smooth_l1_loss_base(bbox_pred,
                                 bbox_targets,
                                 sigma=sigma)
    value = tf.reduce_sum(value, 2)
    value = tf.reshape(value, [-1, 1])

    normalizer = tf.stop_gradient(tf.where(tf.equal(anchor_state, 1)))
    normalizer = tf.to_float(tf.shape(normalizer)[0])
    normalizer = tf.maximum(1.0, normalizer)

    bbox_loss = tf.reduce_sum(
        tf.reduce_sum(value, 1)*outside_mask) / normalizer

    return bbox_loss
//  [python]
def smooth_l1_loss_rcnn(bbox_targets, bbox_pred, anchor_state, sigma=3.0):

    outside_mask = tf.stop_gradient(tf.to_float(tf.greater(anchor_state, 0)))

    bbox_pred = tf.reshape(bbox_pred, [-1, 1, 4])
    bbox_targets = tf.reshape(bbox_targets, [-1, 1, 4])

    value = _smooth_l1_loss_base(bbox_pred,
                                 bbox_targets,
                                 sigma=sigma)
    value = tf.reduce_sum(value, 2)
    value = tf.reshape(value, [-1, 1])

    normalizer = tf.stop_gradient(tf.where(tf.equal(anchor_state, 1)))
    normalizer = tf.cast(tf.shape(normalizer)[0], tf.float32)
    normalizer = tf.maximum(1.0, normalizer)

    bbox_loss = tf.reduce_sum(
        tf.reduce_sum(value, 1)*outside_mask) / normalizer

    return bbox_loss
//  [python]
def call(self, inputs, step_type, network_state=(), training=False):
    observation, action = inputs
    observation_spec, _ = self.input_tensor_spec
    num_outer_dims = nest_utils.get_outer_rank(observation, observation_spec)
    if num_outer_dims not in (1, 2):
      raise ValueError(
          'Input observation must have a batch or batch x time outer shape.'
      )

    has_time_dim = num_outer_dims == 2
    if not has_time_dim:
      # Add a time dimension to the inputs.
      observation = tf.nest.map_structure(
          lambda t: tf.expand_dims(t, 1), observation
      )
      action = tf.nest.map_structure(lambda t: tf.expand_dims(t, 1), action)
      step_type = tf.nest.map_structure(
          lambda t: tf.expand_dims(t, 1), step_type
      )

    observation = t.cast(tf.nest.flatten(observation)[0], tf.float32)
    action = t.cast(tf.nest.flatten(action)[0], tf.float32)

    batch_squash = utils.BatchSquash(2)  # Squash B, and T dims.
    observation = batch_squash.flatten(observation)  # [B, T, ...] -> [BxT, ...]
    action = batch_squash.flatten(action)

    for layer in self._observation_layers:
      observation = layer(observation, training=training)

    for layer in self._action_layers:
      action = layer(action, training=training)

    joint = tf.concat([observation, action], -1)
    for layer in self._joint_layers:
      joint = layer(joint, training=training)

    joint = batch_squash.unflatten(joint)  # [B x T, ...] -> [B, T, ...]

    network_kwargs = {}
    if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):
      network_kwargs['reset_mask'] = tf.equal(
          step_type, time_step.StepType.FIRST, name='mask'
      )

    # Unroll over the time sequence.
    output = self._lstm_network(
        inputs=joint,
        initial_state=network_state,
        training=training,
        **network_kwargs
    )
    if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):
      joint, network_state = output
    else:
      joint = output[0]
      network_state = tf.nest.pack_sequence_as(
          self._lstm_network.cell.state_size, tf.nest.flatten(output[1:])
      )

    output = batch_squash.flatten(joint)  # [B, T, ...] -> [B x T, ...]

    for layer in self._output_layers:
      output = layer(output, training=training)

    q_value = tf.reshape(output, [-1])
    q_value = batch_squash.unflatten(q_value)  # [B x T, ...] -> [B, T, ...]
    if not has_time_dim:
      q_value = tf.squeeze(q_value, axis=1)

    return q_value, network_state
//  [python]
def call(self, inputs, step_type, network_state=(), training=False):
    observation, action = inputs
    observation_spec, _ = self.input_tensor_spec
    num_outer_dims = nest_utils.get_outer_rank(observation, observation_spec)
    if num_outer_dims not in (1, 2):
      raise ValueError(
          'Input observation must have a batch or batch x time outer shape.'
      )

    has_time_dim = num_outer_dims == 2
    if not has_time_dim:
      # Add a time dimension to the inputs.
      observation = tf.nest.map_structure(
          lambda t: tf.expand_dims(t, 1), observation
      )
      action = tf.nest.map_structure(lambda t: tf.expand_dims(t, 1), action)
      step_type = tf.nest.map_structure(
          lambda t: tf.expand_dims(t, 1), step_type
      )

    observation = tf.cast(tf.nest.flatten(observation)[0], tf.float32)
    action = tf.cast(tf.nest.flatten(action)[0], tf.float32)

    batch_squash = utils.BatchSquash(2)  # Squash B, and T dims.
    observation = batch_squash.flatten(observation)  # [B, T, ...] -> [BxT, ...]
    action = batch_squash.flatten(action)

    for layer in self._observation_layers:
      observation = layer(observation, training=training)

    for layer in self._action_layers:
      action = layer(action, training=training)

    joint = tf.concat([observation, action], -1)
    for layer in self._joint_layers:
      joint = layer(joint, training=training)

    joint = batch_squash.unflatten(joint)  # [B x T, ...] -> [B, T, ...]

    network_kwargs = {}
    if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):
      network_kwargs['reset_mask'] = tf.equal(
          step_type, time_step.StepType.FIRST, name='mask'
      )

    # Unroll over the time sequence.
    output = self._lstm_network(
        inputs=joint,
        initial_state=network_state,
        training=training,
        **network_kwargs
    )
    if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):
      joint, network_state = output
    else:
      joint = output[0]
      network_state = tf.nest.pack_sequence_as(
          self._lstm_network.cell.state_size, tf.nest.flatten(output[1:])
      )

    output = batch_squash.flatten(joint)  # [B, T, ...] -> [B x T, ...]

    for layer in self._output_layers:
      output = layer(output, training=training)

    q_value = tf.reshape(output, [-1])
    q_value = batch_squash.unflatten(q_value)  # [B x T, ...] -> [B, T, ...]
    if not has_time_dim:
      q_value = tf.squeeze(q_value, axis=1)

    return q_value, network_state
//  [python]
def get_loss_fn(num_classes):
  """Gets the classification loss function."""

  def classification_loss_fn(labels, logits):
    """Classification loss."""
    labels = tf.reshape(labels, [-1])
    log_probs = tf.nn.log_softmax(logits, axis=-1)
    one_hot_labels = tf.one_hot(
        tf.to_int32(labels), depth=num_classes, dtype=tf.float32)
    per_example_loss = -tf.reduce_sum(
        tf.to_float(one_hot_labels) * log_probs, axis=-1)
    return tf.reduce_mean(per_example_loss)

  return classification_loss_fn
//  [python]
def get_loss_fn(num_classes):
  """Gets the classification loss function."""

  def classification_loss_fn(labels, logits):
    """Classification loss."""
    labels = tf.reshape(labels, [-1])
    log_probs = tf.nn.log_softmax(logits, axis=-1)
    one_hot_labels = tf.one_hot(
        tf.cast(labels, dtype=tf.int32), depth=num_classes, dtype=tf.float32)
    per_example_loss = -tf.reduce_sum(
        tf.cast(one_hot_labels, dtype=tf.float32) * log_probs, axis=-1)
    return tf.reduce_mean(per_example_loss)

  return classification_loss_fn
//  [python]
def _tie_breaking_search(self, anchors, mask, boxes, classes):
    """After search, link each anchor ind to the correct map in ground truth."""
    mask = tf.to_float(tf.reshape(mask, [1, 1, 1, -1]))
    anchors = tf.expand_dims(anchors, axis=-1)
    viable = tf.where(tf.squeeze(anchors == mask, axis=0))

    gather_id, _, anchor_id = tf.split(viable, 3, axis=-1)

    boxes = tf.gather_nd(boxes, gather_id)
    classes = tf.gather_nd(classes, gather_id)

    classes = tf.expand_dims(classes, axis=-1)
    classes = tf.to_float(classes)
    anchor_id = tf.to_float(anchor_id)
    return boxes, classes, anchor_id
//  [python]
def _tie_breaking_search(self, anchors, mask, boxes, classes):
    """After search, link each anchor ind to the correct map in ground truth."""
    mask = tf.cast(tf.reshape(mask, [1, 1, 1, -1]), anchors.dtype)
    anchors = tf.expand_dims(anchors, axis=-1)
    viable = tf.where(tf.squeeze(anchors == mask, axis=0))

    gather_id, _, anchor_id = tf.split(viable, 3, axis=-1)

    boxes = tf.gather_nd(boxes, gather_id)
    classes = tf.gather_nd(classes, gather_id)

    classes = tf.expand_dims(classes, axis=-1)
    classes = tf.cast(classes, boxes.dtype)
    anchor_id = tf.cast(anchor_id, boxes.dtype)
    return boxes, classes, anchor_id
//  [python]
def test_reclayer_optimize_out_accum_loop_dyn_size():
    # We want to test for the case where some layer inside the loop
    # generates some dyn size of shape [B] which is different in each loop frame.
    # So outside the loop, the accumulated dyn size should be of shape [T,B] or [B,T].
    # To test this, we first generate some random seq lens based on the input data (shape [B,T,D]).
    from returnn.tf.util.basic import py_print
    from returnn.tf.util.data import batch_dim, Dim

    def _eval_seq_lens(source, **_kwargs):
        # Get some random varying seq lens.
        res = t.cast(4.0 * source(0) / source(1) + 0.3 * t.cast(source(2), t.float32), t.int32) + 1
        res = py_print(res, ["seq lens", res, "step :i", source(2)])
        return res

    new_time_dim = SpatialDim("T_new")
    feat_dim = FeatureDim("F", dimension=13)
    check_reclayer_optimize_out(
        feat_dim=feat_dim,
        subnet_layer_dict={"class": "linear", "from": "combine", "activation": None, "n_out": 3},
        other_subnet_layers={
            "exp_data": {"class": "activation", "from": "data:source", "activation": "exp"},  # >0
            "sum_exp_data": {"class": "reduce", "mode": "sum", "from": "exp_data", "axis": "F"},  # [B]
            "seq_lens": {
                "class": "eval",
                "from": ["sum_exp_data", "base:max_sum_exp_data", ":i"],
                "out_type": {"dtype": "int32"},
                "out_shape": {batch_dim},
                "eval": _eval_seq_lens,
            },  # [B]
            "range": {"class": "range_from_length", "from": "seq_lens", "out_spatial_dim": new_time_dim},  # [T_new]
            "combine": {
                "class": "eval",
                "from": ["data:source", "range"],
                "eval": "source(0) + 0.1 * t.cast(source(1), t.float32)",
                "out_shape": {batch_dim, new_time_dim, feat_dim},
            },  # [B,T_new,D]
        },
        shared_base_net={
            "exp_data": {"class": "activation", "from": "data", "activation": "exp"},  # >0
            "sum_exp_data": {"class": "reduce", "mode": "sum", "from": "exp_data", "axis": "F"},  # [B,T]
            "max_sum_exp_data": {
                "class": "reduce",
                "mode": "max",
                "from": "sum_exp_data",
                "axis": "T",
                "is_output_layer": True,
            },  # [B]
        },
    )
//  [python]
def test_reclayer_optimize_out_accum_loop_dyn_size():
    # We want to test for the case where some layer inside the loop
    # generates some dyn size of shape [B] which is different in each loop frame.
    # So outside the loop, the accumulated dyn size should be of shape [T,B] or [B,T].
    # To test this, we first generate some random seq lens based on the input data (shape [B,T,D]).
    from returnn.tf.util.basic import py_print
    from returnn.tf.util.data import batch_dim, Dim

    def _eval_seq_lens(source, **_kwargs):
        # Get some random varying seq lens.
        res = tf.cast(4.0 * source(0) / source(1) + 0.3 * tf.cast(source(2), tf.float32), tf.int32) + 1
        res = py_print(res, ["seq lens", res, "step :i", source(2)])
        return res

    new_time_dim = SpatialDim("T_new")
    feat_dim = FeatureDim("F", dimension=13)
    check_reclayer_optimize_out(
        feat_dim=feat_dim,
        subnet_layer_dict={"class": "linear", "from": "combine", "activation": None, "n_out": 3},
        other_subnet_layers={
            "exp_data": {"class": "activation", "from": "data:source", "activation": "exp"},  # >0
            "sum_exp_data": {"class": "reduce", "mode": "sum", "from": "exp_data", "axis": "F"},  # [B]
            "seq_lens": {
                "class": "eval",
                "from": ["sum_exp_data", "base:max_sum_exp_data", ":i"],
                "out_type": {"dtype": "int32"},
                "out_shape": {batch_dim},
                "eval": _eval_seq_lens,
            },  # [B]
            "range": {"class": "range_from_length", "from": "seq_lens", "out_spatial_dim": new_time_dim},  # [T_new]
            "combine": {
                "class": "eval",
                "from": ["data:source", "range"],
                "eval": "source(0) + 0.1 * tf.cast(source(1), tf.float32)",
                "out_shape": {batch_dim, new_time_dim, feat_dim},
            },  # [B,T_new,D]
        },
        shared_base_net={
            "exp_data": {"class": "activation", "from": "data", "activation": "exp"},  # >0
            "sum_exp_data": {"class": "reduce", "mode": "sum", "from": "exp_data", "axis": "F"},  # [B,T]
            "max_sum_exp_data": {
                "class": "reduce",
                "mode": "max",
                "from": "sum_exp_data",
                "axis": "T",
                "is_output_layer": True,
            },  # [B]
        },
    )
//  [python]
def _testGeluPerTensorPostop(self, api_mode):

    with ops.name_scope("test"):
      x_f32 = constant_op.constant([[1.51, 1.13, -0.66],
                                    [0.79, 0.98, 1.88]])
      x_min = tf.math.reduce_min(x_f32)
      x_max = tf.math.reduce_max(x_f32)
      x_int8, x_min, x_max = array_ops.quantize(x_f32, x_min, x_max, T=dtypes.quint8, mode="MIN_FIRST", narrow_range=False)

      y_f32 = constant_op.constant([[-1.71, 3.05, 2.05, 0.64],
                                    [ 0.73, 0.23, 1.42, 3.57],
                                    [-0.61, 2.24, 2.69, 2.90]])
      y_min = tf.math.reduce_min(y_f32)
      y_max = tf.math.reduce_max(y_f32)
      y_int8, y_min, y_max = array_ops.quantize(y_f32, y_min, y_max, T=dtypes.qint8, mode="SCALED", round_mode="HALF_TO_EVEN", narrow_range=True)
      
      bias_f32 = constant_op.constant([-1.42, 2.05, -1.00, 0.94], dtype=dtypes.float32)
      bias_int32 = constant_op.constant([420, 50, -10, 940], dtype=dtypes.int32)
      bias_qint32 = tf.cast(bias_int32, dtypes.qint32)

      if api_mode == "V1":
        matmul_int8 = load_ops_library._QuantizedFusedMatMulAndDequantize(a=x_int8, b=y_int8, args=[bias_qint32], min_a=x_min, max_a=x_max,
                                                                          min_b=y_min, max_b=y_max,
                                                                          Toutput=tf.float32, fused_ops=["BiasAdd", "GeluApproximate"],
                                                                          input_quant_mode="MIN_FIRST") 

        matmul_int8 = array_ops.identity(matmul_int8)

      elif api_mode == "V2":
        matmul_int8 = load_ops_library._QuantizedMatMul(device_inputs=[x_int8, y_int8, bias_qint32], 
                                                                       host_inputs=[x_min, x_max, y_min, y_max],
                                                                       T1=dtypes.quint8, T2=dtypes.qint8, Tbias=dtypes.qint32,
                                                                       Tdevice_outputs=[dtypes.float32],
                                                                       fused_ops=["BiasAdd", "GeluApproximate", "Dequantize"],
                                                                       input_quant_mode="MIN_FIRST") 

        matmul_int8 = array_ops.identity(matmul_int8[0])[0]
      
      matmul_f32 = itex.ops.gelu(tf.nn.bias_add(tf.linalg.matmul(x_f32, y_f32), bias_f32), approximate=True)
      matmul_f32 = array_ops.identity(matmul_f32)

      with self.session(use_gpu=True) as sess:
        # Run twice to check the funtionality of object cache
        for _ in range(2):
          matmul_int8_res = sess.run(matmul_int8)
          matmul_f32_res = sess.run(matmul_f32)
//  [python]
def build_label_per_path(self,
                           key,
                           boxes,
                           classes,
                           width,
                           height,
                           iou_index=None):
    """Builds the labels for one path."""
    stride = self.level_strides[key]
    scale_xy = self.center_radius[key] if self.center_radius is not None else 1

    width = t.cast(width // stride, boxes.dtype)
    height = t.cast(height // stride, boxes.dtype)

    if self.anchor_free_level_limits is None:
      (boxes, classes, anchors, num_anchors) = self._get_anchor_id(
          key, boxes, classes, width, height, stride, iou_index=iou_index)
      boxes, classes, centers = self._get_centers(boxes, classes, anchors,
                                                  width, height, scale_xy)
      ind_mask = t.ones_like(classes)
      updates = t.concat([boxes, ind_mask, classes], axis=-1)
    else:
      num_anchors = 1
      (centers, updates) = self._get_anchor_free(key, boxes, classes, height,
                                                 width, stride, scale_xy)
      boxes, ind_mask, classes = t.split(updates, [4, 1, 1], axis=-1)

    width = t.cast(width, t.int32)
    height = t.cast(height, t.int32)
    full = t.zeros([height, width, num_anchors, 1], dtype=classes.dtype)
    full = t.tensor_scatter_nd_add(full, centers, ind_mask)

    num_instances = int(self.num_instances[key])
    centers = preprocessing_ops.pad_max_instances(
        centers, num_instances, pad_value=0, pad_axis=0)
    updates = preprocessing_ops.pad_max_instances(
        updates, num_instances, pad_value=0, pad_axis=0)

    updates = t.cast(updates, self.dtype)
    full = t.cast(full, self.dtype)
    return centers, updates, full
//  [python]
def build_label_per_path(self,
                           key,
                           boxes,
                           classes,
                           width,
                           height,
                           iou_index=None):
    """Builds the labels for one path."""
    stride = self.level_strides[key]
    scale_xy = self.center_radius[key] if self.center_radius is not None else 1

    width = tf.cast(width // stride, boxes.dtype)
    height = tf.cast(height // stride, boxes.dtype)

    if self.anchor_free_level_limits is None:
      (boxes, classes, anchors, num_anchors) = self._get_anchor_id(
          key, boxes, classes, width, height, stride, iou_index=iou_index)
      boxes, classes, centers = self._get_centers(boxes, classes, anchors,
                                                  width, height, scale_xy)
      ind_mask = tf.ones_like(classes)
      updates = tf.concat([boxes, ind_mask, classes], axis=-1)
    else:
      num_anchors = 1
      (centers, updates) = self._get_anchor_free(key, boxes, classes, height,
                                                 width, stride, scale_xy)
      boxes, ind_mask, classes = tf.split(updates, [4, 1, 1], axis=-1)

    width = tf.cast(width, tf.int32)
    height = tf.cast(height, tf.int32)
    full = tf.zeros([height, width, num_anchors, 1], dtype=classes.dtype)
    full = tf.tensor_scatter_nd_add(full, centers, ind_mask)

    num_instances = int(self.num_instances[key])
    centers = preprocessing_ops.pad_max_instances(
        centers, num_instances, pad_value=0, pad_axis=0)
    updates = preprocessing_ops.pad_max_instances(
        updates, num_instances, pad_value=0, pad_axis=0)

    updates = tf.cast(updates, self.dtype)
    full = tf.cast(full, self.dtype)
    return centers, updates, full
//  [python]
def augment_images(
        self, images, transformations, resize_method="bilinear", **kwargs
    ):
        # unpackage augmentation arguments
        scaled_sizes = transformations["scaled_sizes"]
        offsets = transformations["offsets"]
        inputs_for_resize_and_crop_single_image = {
            "images": images,
            "scaled_sizes": scaled_sizes,
            "offsets": offsets,
        }
        scaled_images = tf.map_fn(
            lambda x: self.resize_and_crop_single_image(
                x, resize_method=resize_method
            ),
            inputs_for_resize_and_crop_single_image,
            fn_output_signature=tf.float32,
        )
        return tf.cast(scaled_images, self.compute_dtype)
//  [python]
def __call__(self, boxes, classes, width, height):
    """Builds the labels for a single image, not functional in batch mode.

    Args:
      boxes: `Tensor` of shape [None, 4] indicating the object locations in an
        image.
      classes: `Tensor` of shape [None] indicating the each objects classes.
      width: `int` for the images width.
      height: `int` for the images height.

    Returns:
      centers: `Tensor` of shape [None, 3] of indexes in the final grid where
        boxes are located.
      updates: `Tensor` of shape [None, 8] the value to place in the final grid.
      full: `Tensor` of [width/stride, height/stride, num_anchors, 1] holding
        a mask of where boxes are locates for confidence losses.
    """
    indexes = {}
    updates = {}
    true_grids = {}
    iou_index = None

    boxes = box_ops.yxyx_to_xcycwh(boxes)
    if not self.best_matches_only and self.anchor_free_level_limits is None:
      # stitch and search boxes across fpn levels
      anchorsvec = []
      for stitch in self.anchors:
        anchorsvec.extend(self.anchors[stitch])

      stride = t.cast([width, height], boxes.dtype)
      # get the best anchor for each box
      iou_index, _ = get_best_anchor(
          boxes,
          anchorsvec,
          stride,
          width=1.0,
          height=1.0,
          best_match_only=False,
          use_tie_breaker=self.use_tie_breaker,
          iou_thresh=self.match_threshold)

    for key in self.keys:
      indexes[key], updates[key], true_grids[key] = self.build_label_per_path(
          key, boxes, classes, width, height, iou_index=iou_index)
    return indexes, updates, true_grids
//  [python]
def __call__(self, boxes, classes, width, height):
    """Builds the labels for a single image, not functional in batch mode.

    Args:
      boxes: `Tensor` of shape [None, 4] indicating the object locations in an
        image.
      classes: `Tensor` of shape [None] indicating the each objects classes.
      width: `int` for the images width.
      height: `int` for the images height.

    Returns:
      centers: `Tensor` of shape [None, 3] of indexes in the final grid where
        boxes are located.
      updates: `Tensor` of shape [None, 8] the value to place in the final grid.
      full: `Tensor` of [width/stride, height/stride, num_anchors, 1] holding
        a mask of where boxes are locates for confidence losses.
    """
    indexes = {}
    updates = {}
    true_grids = {}
    iou_index = None

    boxes = box_ops.yxyx_to_xcycwh(boxes)
    if not self.best_matches_only and self.anchor_free_level_limits is None:
      # stitch and search boxes across fpn levels
      anchorsvec = []
      for stitch in self.anchors:
        anchorsvec.extend(self.anchors[stitch])

      stride = tf.cast([width, height], boxes.dtype)
      # get the best anchor for each box
      iou_index, _ = get_best_anchor(
          boxes,
          anchorsvec,
          stride,
          width=1.0,
          height=1.0,
          best_match_only=False,
          use_tie_breaker=self.use_tie_breaker,
          iou_thresh=self.match_threshold)

    for key in self.keys:
      indexes[key], updates[key], true_grids[key] = self.build_label_per_path(
          key, boxes, classes, width, height, iou_index=iou_index)
    return indexes, updates, true_grids
//  [python]
def get_random_transformation_batch(
        self, batch_size, images=None, **kwargs
    ):
        heights, widths = self._get_image_shape(images)
        image_shapes = t.concat((heights, widths), axis=-1)

        scaled_sizes = t.round(
            image_shapes * self.scale_factor(shape=(batch_size, 1))
        )
        scales = t.where(
            t.less(
                scaled_sizes[..., 0] / image_shapes[..., 0],
                scaled_sizes[..., 1] / image_shapes[..., 1],
            ),
            scaled_sizes[..., 0] / image_shapes[..., 0],
            scaled_sizes[..., 1] / image_shapes[..., 1],
        )

        scaled_sizes = t.round(image_shapes * scales[..., t.newaxis])
        image_scales = scaled_sizes / image_shapes

        max_offsets = scaled_sizes - self.crop_size
        max_offsets = t.where(
            t.less(max_offsets, 0), t.zeros_like(max_offsets), max_offsets
        )
        offsets = max_offsets * self._random_generator.uniform(
            shape=(batch_size, 2), minval=0, maxval=1, dtype=t.float32
        )
        offsets = t.cast(offsets, t.int32)
        return {
            "image_scales": image_scales,
            "scaled_sizes": scaled_sizes,
            "offsets": offsets,
        }
//  [python]
def get_random_transformation_batch(
        self, batch_size, images=None, **kwargs
    ):
        heights, widths = self._get_image_shape(images)
        image_shapes = tf.cast(
            tf.concat((heights, widths), axis=-1), dtype=tf.float32
        )

        scaled_sizes = tf.round(
            image_shapes * self.scale_factor(shape=(batch_size, 1))
        )
        scales = tf.where(
            tf.less(
                scaled_sizes[..., 0] / image_shapes[..., 0],
                scaled_sizes[..., 1] / image_shapes[..., 1],
            ),
            scaled_sizes[..., 0] / image_shapes[..., 0],
            scaled_sizes[..., 1] / image_shapes[..., 1],
        )

        scaled_sizes = tf.round(image_shapes * scales[..., tf.newaxis])
        image_scales = scaled_sizes / image_shapes

        max_offsets = scaled_sizes - self.crop_size
        max_offsets = tf.where(
            tf.less(max_offsets, 0), tf.zeros_like(max_offsets), max_offsets
        )
        offsets = max_offsets * self._random_generator.uniform(
            shape=(batch_size, 2), minval=0, maxval=1, dtype=tf.float32
        )
        offsets = tf.cast(offsets, tf.int32)
        return {
            "image_scales": image_scales,
            "scaled_sizes": scaled_sizes,
            "offsets": offsets,
        }
//  [python]
def _get_centers(self, boxes, classes, anchors, width, height, scale_xy):
    """Find the object center assignments in an anchor based paradigm."""
    offset = t.cast(0.5 * (scale_xy - 1), boxes.dtype)

    grid_xy, _ = t.split(boxes, 2, axis=-1)
    wh_scale = t.cast(t.convert_to_tensor([width, height]), boxes.dtype)

    grid_xy = grid_xy * wh_scale
    centers = t.math.floor(grid_xy)

    if offset != 0.0:
      clamp = lambda x, ma: t.maximum(  # pylint:disable=g-long-lambda
          t.minimum(x, t.cast(ma, x.dtype)), t.zeros_like(x))

      grid_xy_index = grid_xy - centers
      positive_shift = ((grid_xy_index < offset) & (grid_xy > 1.))
      negative_shift = ((grid_xy_index > (1 - offset)) & (grid_xy <
                                                          (wh_scale - 1.)))

      zero, _ = t.split(t.ones_like(positive_shift), 2, axis=-1)
      shift_mask = t.concat([zero, positive_shift, negative_shift], axis=-1)
      offset = t.cast([[0, 0], [1, 0], [0, 1], [-1, 0], [0, -1]],
                       offset.dtype) * offset

      num_shifts = t.shape(shift_mask)
      num_shifts = num_shifts[-1]
      boxes = t.tile(t.expand_dims(boxes, axis=-2), [1, num_shifts, 1])
      classes = t.tile(t.expand_dims(classes, axis=-2), [1, num_shifts, 1])
      anchors = t.tile(t.expand_dims(anchors, axis=-2), [1, num_shifts, 1])

      shift_mask = t.cast(shift_mask, boxes.dtype)
      shift_ind = shift_mask * t.range(0, num_shifts, dtype=boxes.dtype)
      shift_ind = shift_ind - (1 - shift_mask)
      shift_ind = t.expand_dims(shift_ind, axis=-1)

      boxes_and_centers = t.concat([boxes, classes, anchors, shift_ind],
                                    axis=-1)
      boxes_and_centers = t.reshape(boxes_and_centers, [-1, 7])
      _, center_ids = t.split(boxes_and_centers, [6, 1], axis=-1)

      select = t.where(center_ids >= 0)
      select, _ = t.split(select, 2, axis=-1)

      boxes_and_centers = t.gather_nd(boxes_and_centers, select)

      center_ids = t.gather_nd(center_ids, select)
      center_ids = t.cast(center_ids, t.int32)
      shifts = t.gather_nd(offset, center_ids)

      boxes, classes, anchors, _ = t.split(
          boxes_and_centers, [4, 1, 1, 1], axis=-1)
      grid_xy, _ = t.split(boxes, 2, axis=-1)
      centers = t.math.floor(grid_xy * wh_scale - shifts)
      centers = clamp(centers, wh_scale - 1)

    x, y = t.split(centers, 2, axis=-1)
    centers = t.cast(t.concat([y, x, anchors], axis=-1), t.int32)
    return boxes, classes, centers
//  [python]
def _get_centers(self, boxes, classes, anchors, width, height, scale_xy):
    """Find the object center assignments in an anchor based paradigm."""
    offset = tf.cast(0.5 * (scale_xy - 1), boxes.dtype)

    grid_xy, _ = tf.split(boxes, 2, axis=-1)
    wh_scale = tf.cast(tf.convert_to_tensor([width, height]), boxes.dtype)

    grid_xy = grid_xy * wh_scale
    centers = tf.math.floor(grid_xy)

    if offset != 0.0:
      clamp = lambda x, ma: tf.maximum(  # pylint:disable=g-long-lambda
          tf.minimum(x, tf.cast(ma, x.dtype)), tf.zeros_like(x))

      grid_xy_index = grid_xy - centers
      positive_shift = ((grid_xy_index < offset) & (grid_xy > 1.))
      negative_shift = ((grid_xy_index > (1 - offset)) & (grid_xy <
                                                          (wh_scale - 1.)))

      zero, _ = tf.split(tf.ones_like(positive_shift), 2, axis=-1)
      shift_mask = tf.concat([zero, positive_shift, negative_shift], axis=-1)
      offset = tf.cast([[0, 0], [1, 0], [0, 1], [-1, 0], [0, -1]],
                       offset.dtype) * offset

      num_shifts = tf.shape(shift_mask)
      num_shifts = num_shifts[-1]
      boxes = tf.tile(tf.expand_dims(boxes, axis=-2), [1, num_shifts, 1])
      classes = tf.tile(tf.expand_dims(classes, axis=-2), [1, num_shifts, 1])
      anchors = tf.tile(tf.expand_dims(anchors, axis=-2), [1, num_shifts, 1])

      shift_mask = tf.cast(shift_mask, boxes.dtype)
      shift_ind = shift_mask * tf.range(0, num_shifts, dtype=boxes.dtype)
      shift_ind = shift_ind - (1 - shift_mask)
      shift_ind = tf.expand_dims(shift_ind, axis=-1)

      boxes_and_centers = tf.concat([boxes, classes, anchors, shift_ind],
                                    axis=-1)
      boxes_and_centers = tf.reshape(boxes_and_centers, [-1, 7])
      _, center_ids = tf.split(boxes_and_centers, [6, 1], axis=-1)

      select = tf.where(center_ids >= 0)
      select, _ = tf.split(select, 2, axis=-1)

      boxes_and_centers = tf.gather_nd(boxes_and_centers, select)

      center_ids = tf.gather_nd(center_ids, select)
      center_ids = tf.cast(center_ids, tf.int32)
      shifts = tf.gather_nd(offset, center_ids)

      boxes, classes, anchors, _ = tf.split(
          boxes_and_centers, [4, 1, 1, 1], axis=-1)
      grid_xy, _ = tf.split(boxes, 2, axis=-1)
      centers = tf.math.floor(grid_xy * wh_scale - shifts)
      centers = clamp(centers, wh_scale - 1)

    x, y = tf.split(centers, 2, axis=-1)
    centers = tf.cast(tf.concat([y, x, anchors], axis=-1), tf.int32)
    return boxes, classes, centers
//  [python]
def _get_image_shape(self, images):
        if isinstance(images, tf.RaggedTensor):
            heights = tf.reshape(images.row_lengths(), (-1, 1))
            widths = tf.reshape(
                tf.reduce_max(images.row_lengths(axis=2), 1), (-1, 1)
            )
        else:
            batch_size = tf.shape(images)[0]
            heights = tf.repeat(tf.shape(images)[H_AXIS], repeats=[batch_size])
            heights = tf.reshape(heights, shape=(-1, 1))
            widths = tf.repeat(tf.shape(images)[W_AXIS], repeats=[batch_size])
            widths = tf.reshape(widths, shape=(-1, 1))
        return tf.to_int32(heights), tf.to_int32(widths)
//  [python]
def _get_image_shape(self, images):
        if isinstance(images, tf.RaggedTensor):
            heights = tf.reshape(images.row_lengths(), (-1, 1))
            widths = tf.reshape(
                tf.reduce_max(images.row_lengths(axis=2), 1), (-1, 1)
            )
        else:
            batch_size = tf.shape(images)[0]
            heights = tf.repeat(tf.shape(images)[H_AXIS], repeats=[batch_size])
            heights = tf.reshape(heights, shape=(-1, 1))
            widths = tf.repeat(tf.shape(images)[W_AXIS], repeats=[batch_size])
            widths = tf.reshape(widths, shape=(-1, 1))
        return tf.cast(heights, dtype=tf.int32), tf.cast(widths, dtype=tf.int32)