jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _suppression_loop_body(boxes, iou_threshold, output_size, idx):
  """Process boxes in the range [idx*_NMS_TILE_SIZE, (idx+1)*_NMS_TILE_SIZE).

  Args:
    boxes: a tensor with a shape of [1, anchors, 4].
    iou_threshold: a float representing the threshold for deciding whether boxes
      overlap too much with respect to IOU.
    output_size: an int32 tensor of size [1]. Representing the number of
      selected boxes.
    idx: an integer scalar representing induction variable.

  Returns:
    boxes: updated boxes.
    iou_threshold: pass down iou_threshold to the next iteration.
    output_size: the updated output_size.
    idx: the updated induction variable.
  """
  num_tiles = tf.shape(boxes)[1] // _NMS_TILE_SIZE

  # Iterates over tiles that can possibly suppress the current tile.
  box_slice = tf.slice(boxes, [0, idx * _NMS_TILE_SIZE, 0],
                       [1, _NMS_TILE_SIZE, 4])
  _, box_slice, _, _ = tf.while_loop(
      lambda _boxes, _box_slice, _threshold, inner_idx: inner_idx < idx,
      _cross_suppression, [boxes, box_slice, iou_threshold,
                           tf.constant(0)])

  # Iterates over the current tile to compute self-suppression.
  iou = batch_iou(box_slice, box_slice)
  mask = tf.expand_dims(
      tf.reshape(tf.range(_NMS_TILE_SIZE), [1, -1]) > tf.reshape(
          tf.range(_NMS_TILE_SIZE), [-1, 1]), 0)
  iou *= tf.cast(tf.logical_and(mask, iou >= iou_threshold), iou.dtype)
  suppressed_iou, _, _, _ = tf.while_loop(
      lambda _iou, _threshold, loop_condition, _iou_sum: loop_condition,
      _self_suppression,
      [iou, iou_threshold,
       tf.constant(True),
       tf.reduce_sum(iou, [1, 2])])
  suppressed_box = tf.reduce_sum(suppressed_iou, 1) > 0
  box_slice *= tf.expand_dims(1.0 - tf.cast(suppressed_box, box_slice.dtype), 2)

  # Uses box_slice to update the input boxes.
  mask = tf.reshape(
      tf.cast(tf.equal(tf.range(num_tiles), idx), boxes.dtype), [1, -1, 1, 1])
  boxes = tf.tile(tf.expand_dims(box_slice, [1]),
                  [1, num_tiles, 1, 1]) * mask + tf.reshape(
                      boxes, [1, num_tiles, _NMS_TILE_SIZE, 4]) * (1 - mask)
  boxes = tf.reshape(boxes, [1, -1, 4])

  # Updates output_size.
  output_size += tf.reduce_sum(
      tf.cast(tf.reduce_any(box_slice > 0, [2]), tf.int32), [1])
  return boxes, iou_threshold, output_size, idx + 1
//  [python]
def _compute_stacked_offsets(sizes, repeats):
  """Computes offsets to add to indices of stacked tensors (Tensorflow).

  When a set of tensors are stacked, the indices of those from the second on
  must be offset in order to be able to index into the stacked tensor. This
  computes those offsets.

  Args:
    sizes: A 1D `Tensor` of the sizes per graph.
    repeats: A 1D `Tensor` of the number of repeats per graph.

  Returns:
    A 1D `Tensor` containing the index offset per graph.
  """
  sizes = tf.convert_to_tensor(sizes[:-1])
  offset_values = tf.cumsum(tf.concat([[0], sizes], 0))
  return repeat(offset_values, repeats)
//  [python]
def _compute_stacked_offsets(sizes, repeats):
  """Computes offsets to add to indices of stacked tensors (Tensorflow).

  When a set of tensors are stacked, the indices of those from the second on
  must be offset in order to be able to index into the stacked tensor. This
  computes those offsets.

  Args:
    sizes: A 1D `Tensor` of the sizes per graph.
    repeats: A 1D `Tensor` of the number of repeats per graph.

  Returns:
    A 1D `Tensor` containing the index offset per graph.
  """
  sizes = tf.cast(tf.convert_to_tensor(sizes[:-1]), tf.int32)
  offset_values = tf.cumsum(tf.concat([[0], sizes], 0))
  return repeat(offset_values, repeats)
//  [python]
def warmup_lr(init_lr, global_step, warmup_step, num_gpu):
    def warmup(end_lr, global_step, warmup_step):
        start_lr = end_lr * 0.1
        global_step = tf.cast(global_step, tf.float32)
        return start_lr + (end_lr - start_lr) * global_step / warmup_step

    def decay(start_lr, global_step, num_gpu):
        lr = tf.train.piecewise_constant(global_step,
                                         boundaries=[np.int64(cfgs.DECAY_STEP[0] // num_gpu),
                                                     np.int64(cfgs.DECAY_STEP[1] // num_gpu),
                                                     np.int64(cfgs.DECAY_STEP[2] // num_gpu)],
                                         values=[start_lr, start_lr / 10., start_lr / 100., start_lr / 1000.])
        return lr

    return tf.cond(tf.less_equal(global_step, warmup_step),
                   true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                   false_fn=lambda: decay(init_lr, global_step, num_gpu))
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        causal_attention_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        """Input shape: Batch x Time x Channel"""

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.q_proj(inputs=hidden_states)
        mixed_key_layer = self.k_proj(inputs=hidden_states)
        mixed_value_layer = self.v_proj(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)

        # apply the causal_attention_mask first
        if causal_attention_mask is not None:
            # Apply the causal attention mask (precomputed for all layers in TFCLIPModel call() function)
            attention_scores = tf.add(attention_scores, causal_attention_mask)

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

        # Normalize the attention scores to probabilities.
        _attention_probs = stable_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)

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

        # (batch_size, seq_len_q, embed_dim)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.embed_dim))

        attention_output = self.out_proj(attention_output, training=training)
        # In TFBert, attention weights are returned after dropout.
        # However, in CLIP, they are returned before dropout.
        outputs = (attention_output, _attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def __init__(self, model_dim, warmup_steps=4000):
    super(CustomSchedule, self).__init__()
    self.model_dim = model_dim
    self.model_dim = tf.cast(self.model_dim, tf.float32)
    self.warmup_steps = warmup_steps
//  [python]
def _loss_function(self, y_true, y_pred):
    loss = self._loss_object(y_true, y_pred)
    mask = tf.cast(tf.not_equal(y_true, 0), tf.float32)
    loss = tf.multiply(loss, mask)
    return tf.reduce_mean(loss)
//  [python]
def _create_complete_edges_from_nodes_dynamic(n_node, exclude_self_edges):
  """Creates complete edges for a graph with `n_node`.

  Args:
    n_node: (integer scalar `Tensor`) The number of nodes.
    exclude_self_edges: (bool) Excludes self-connected edges.

  Returns:
    A dict of RECEIVERS, SENDERS and N_EDGE data (`Tensor`s of rank 1).
  """
  rng = tf.range(n_node)
  receivers, senders = tf.meshgrid(rng, rng)
  n_edge = n_node * n_node

  if exclude_self_edges:
    ind = tf.cast(1 - tf.eye(n_node), bool)
    receivers = tf.boolean_mask(receivers, ind)
    senders = tf.boolean_mask(senders, ind)
    n_edge -= n_node

  receivers = tf.reshape(tf.cast(receivers, tf.int32), [n_edge])
  senders = tf.reshape(tf.cast(senders, tf.int32), [n_edge])
  n_edge = tf.reshape(n_edge, [1])

  return {RECEIVERS: receivers, SENDERS: senders, N_EDGE: n_edge}
//  [python]
def get_mask(valid_length, full_length):
  """Returns a mask given the valid length of a vector with trailing padding.

  This is useful for masking out padded elements from a loss. For example
  
  input_graphs_tuple = ...
  input_graphs_tuple_size = get_graphs_tuple_size(input_graphs_tuple)
  padded_input_graphs_tuple = pad_graphs_tuple(input_graphs_tuple,
       pad_nodes_to,...)


  per_node_loss # After graph_network computation.

  nodes_mask = get_mask(
      input_graphs_tuple_size.nodes, pad_nodes_to)

  masked_per_node_loss = per_node_loss * tf.cast(
      nodes_mask, per_node_loss.dtype)
  

  Args:
    valid_length: Length of the valid elements.
    full_length: Full length of the vector after padding.

  Returns:
    Boolean mask of shape [full_length], where all values are set to `True`
    except for the last `padding_length` which are set to False.
  """
  valid_length = t.cast(valid_length, t.int32)
  full_length = t.cast(full_length, t.int32)
  # This implementation allows for statically sized shapes, rather than
  # using concat([tf.ones([valid_length]), tf.zeros([full_length-valid_length])]
  # which has intermediate tensors with shapes not know statically.
  field_mask = tf.range(full_length)
  field_mask = field_mask < valid_length
  return field_mask
//  [python]
def get_mask(valid_length, full_length):
  """Returns a mask given the valid length of a vector with trailing padding.

  This is useful for masking out padded elements from a loss. For example
  ```
  input_graphs_tuple = ...
  input_graphs_tuple_size = get_graphs_tuple_size(input_graphs_tuple)
  padded_input_graphs_tuple = pad_graphs_tuple(input_graphs_tuple,
       pad_nodes_to,...)


  per_node_loss # After graph_network computation.

  nodes_mask = get_mask(
      input_graphs_tuple_size.nodes, pad_nodes_to)

  masked_per_node_loss = per_node_loss * tf.cast(
      nodes_mask, per_node_loss.dtype)
  ```

  Args:
    valid_length: Length of the valid elements.
    full_length: Full length of the vector after padding.

  Returns:
    Boolean mask of shape [full_length], where all values are set to `True`
    except for the last `padding_length` which are set to False.
  """
  valid_length = tf.cast(valid_length, tf.int32)
  full_length = tf.cast(full_length, tf.int32)
  # This implementation allows for statically sized shapes, rather than
  # using concat([tf.ones([valid_length]), tf.zeros([full_length-valid_length])]
  # which has intermediate tensors with shapes not know statically.
  field_mask = tf.range(full_length)
  field_mask = field_mask < valid_length
  return field_mask
//  [python]
def get_random_mask(self):
    sample = tf.random.uniform(shape=self._node_features.shape, maxval=1)
    mask = tf.to_float(
        tf.math.greater(self._probabilities, sample)
    )
    return mask
//  [python]
def get_random_mask(self):
    sample = tf.random.uniform(shape=self._node_features.shape, maxval=1)
    mask = tf.cast(
        tf.math.greater(self._probabilities, sample), dtype=tf.float32
    )
    return mask
//  [python]
def _accuracy_function(self, real, pred):
    accuracies = tf.equal(real, tf.argmax(pred, output_type=tf.int32, axis=2))
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    accuracies = tf.math.logical_and(mask, accuracies)
    accuracies = tf.to_float(accuracies)
    mask = tf.to_float(mask)
    return tf.reduce_sum(accuracies) / tf.reduce_sum(mask)
//  [python]
def _accuracy_function(self, real, pred):
    accuracies = tf.equal(real, tf.argmax(pred, output_type=tf.int32, axis=2))
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    accuracies = tf.math.logical_and(mask, accuracies)
    accuracies = tf.cast(accuracies, dtype=tf.float32)
    mask = tf.cast(mask, dtype=tf.float32)
    return tf.reduce_sum(accuracies) / tf.reduce_sum(mask)
//  [python]
def compute_probability_mask(self):
    nones = tf.math.count_nonzero(self._node_features).numpy()
    nzeros = self._node_features.shape[0] * self._node_features.shape[1] - nones
    probability_zeros = nones / nzeros / self._ones_ratio * self._negative_ratio
    probabilities = tf.math.scalar_mul(
        probability_zeros,
        tf.to_float(tf.equal(self._node_features, 0)),
    ) + tf.math.scalar_mul(
        1 / self._ones_ratio,
        tf.to_float(tf.equal(self._node_features, 1)),
    )
    return probabilities
//  [python]
def compute_probability_mask(self):
    nones = tf.math.count_nonzero(self._node_features).numpy()
    nzeros = self._node_features.shape[0] * self._node_features.shape[1] - nones
    probability_zeros = nones / nzeros / self._ones_ratio * self._negative_ratio
    probabilities = tf.math.scalar_mul(
        probability_zeros,
        tf.cast(tf.equal(self._node_features, 0), dtype=tf.float32),
    ) + tf.math.scalar_mul(
        1 / self._ones_ratio,
        tf.cast(tf.equal(self._node_features, 1), dtype=tf.float32),
    )
    return probabilities
//  [python]
def _decoder(decoder_inputs, encoder_inputs, encoder_outputs, encoder_state,
             target_vocab_size, embedding_dim, num_layers,
             num_hidden_units, kernel_size, dropout_rate):
  """Decoder component."""
  x_decoder = _embedding(decoder_inputs, target_vocab_size, embedding_dim,
                         dropout_rate)

  encoder_padding = tf.expand_dims(tf.math.equal(encoder_inputs, 0), axis=-1)
  real_lengths = tf.cast(tf.shape(encoder_outputs)[1], tf.float32)
  real_lengths -= tf.math.reduce_sum(
      tf.cast(encoder_padding, tf.float32), axis=1, keepdims=True)
  encoder_mask = tf.expand_dims(tf.math.not_equal(encoder_inputs, 0), axis=-1)
  encoder_mask = tf.cast(encoder_mask, tf.float32)
  decoder_mask = tf.cast(tf.expand_dims(tf.math.not_equal(
      decoder_inputs, 0), axis=-1), tf.float32)
  encoder_outputs *= encoder_mask
  encoder_state *= encoder_mask

  attention_scale = real_lengths * tf.math.rsqrt(real_lengths)
  attention_mask = tf.linalg.matmul(decoder_mask, encoder_mask,
                                    transpose_b=True)
  if x_decoder.shape[2] != num_hidden_units:
    x_decoder = tf.keras.layers.Conv1D(num_hidden_units, 1)(x_decoder)
  g = x_decoder
  layer_attention = []
  for i in range(num_layers):
    # Output of the current decoder layer `h_i`.
    h = _layer(x_decoder, _decoder_conv_block, kernel_size,
               num_hidden=num_hidden_units, dilation_rate=2 ** i)
    h *= decoder_mask
    h_residual = h
    # See Eq. (1) in Gehring et. al. (2017).
    d = tf.keras.layers.Dense(num_hidden_units)(h) + g
    d *= _RESIDUAL_SCALING_FACTOR
    d *= decoder_mask
    dz = tf.linalg.matmul(d, tf.transpose(encoder_outputs, [0, 2, 1]))
    a = tf.keras.layers.Softmax()(dz, mask=attention_mask)
    layer_attention.append(a)
    # See Eq. (2) in Gehring et. al. (2017).
    c = tf.linalg.matmul(a, encoder_state)
    c *= attention_scale
    h = (h_residual + c) * _RESIDUAL_SCALING_FACTOR
    h = tf.keras.layers.Dense(num_hidden_units)(h)
    h = tf.keras.layers.Dropout(rate=dropout_rate)(h)
    x_decoder = h

  return x_decoder, layer_attention
//  [python]
def call(self, inputs, outer_rank, training=False, mask=None):
    if inputs.dtype != self._sample_spec.dtype:
      raise ValueError(
          'Inputs to NormalProjectionNetwork must match the sample_spec.dtype.'
      )

    if mask is not None:
      raise NotImplementedError(
          'NormalProjectionNetwork does not yet implement action masking; got '
          'mask={}'.format(mask)
      )

    # outer_rank is needed because the projection is not done on the raw
    # observations so getting the outer rank is hard as there is no spec to
    # compare to.
    batch_squash = network_utils.BatchSquash(outer_rank)
    inputs = batch_squash.flatten(inputs)

    means = self._means_projection_layer(inputs, training=training)
    means = tf.reshape(means, [-1] + self._sample_spec.shape.as_list())

    # If scaling the distribution later, use a normalized mean.
    if not self._scale_distribution and self._mean_transform is not None:
      means = self._mean_transform(means, self._sample_spec)
    means = t.cast(means, self._sample_spec.dtype)

    if self._state_dependent_std:
      stds = self._stddev_projection_layer(inputs, training=training)
    else:
      stds = self._bias(tf.zeros_like(means), training=training)
      stds = tf.reshape(stds, [-1] + self._sample_spec.shape.as_list())

    if self._std_transform is not None:
      stds = self._std_transform(stds)
    stds = t.cast(stds, self._sample_spec.dtype)

    means = batch_squash.unflatten(means)
    stds = batch_squash.unflatten(stds)

    return self.output_spec.build_distribution(loc=means, scale=stds), ()
//  [python]
def call(self, inputs, outer_rank, training=False, mask=None):
    if inputs.dtype != self._sample_spec.dtype:
      raise ValueError(
          'Inputs to NormalProjectionNetwork must match the sample_spec.dtype.'
      )

    if mask is not None:
      raise NotImplementedError(
          'NormalProjectionNetwork does not yet implement action masking; got '
          'mask={}'.format(mask)
      )

    # outer_rank is needed because the projection is not done on the raw
    # observations so getting the outer rank is hard as there is no spec to
    # compare to.
    batch_squash = network_utils.BatchSquash(outer_rank)
    inputs = batch_squash.flatten(inputs)

    means = self._means_projection_layer(inputs, training=training)
    means = tf.reshape(means, [-1] + self._sample_spec.shape.as_list())

    # If scaling the distribution later, use a normalized mean.
    if not self._scale_distribution and self._mean_transform is not None:
      means = self._mean_transform(means, self._sample_spec)
    means = tf.cast(means, self._sample_spec.dtype)

    if self._state_dependent_std:
      stds = self._stddev_projection_layer(inputs, training=training)
    else:
      stds = self._bias(tf.zeros_like(means), training=training)
      stds = tf.reshape(stds, [-1] + self._sample_spec.shape.as_list())

    if self._std_transform is not None:
      stds = self._std_transform(stds)
    stds = tf.cast(stds, self._sample_spec.dtype)

    means = batch_squash.unflatten(means)
    stds = batch_squash.unflatten(stds)

    return self.output_spec.build_distribution(loc=means, scale=stds), ()
//  [python]
def create_position_ids_from_input_ids(self, input_ids, past_key_values_length=0):
        """
        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 = t.cast(tf.math.not_equal(input_ids, self.padding_idx), dtype=input_ids.dtype)
        incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask

        return incremental_indices + self.padding_idx
//  [python]
def create_position_ids_from_input_ids(self, input_ids, past_key_values_length=0):
        """
        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) + past_key_values_length) * mask

        return incremental_indices + self.padding_idx
//  [python]
def call(self, inputs):
    graph_structure = inputs[0]
    node_embeddings = inputs[1]

    number_of_nodes, number_of_features = node_embeddings.shape
    probabilities = tf.repeat(self._feature_mask_rate, number_of_features)
    sample = tf.random.uniform(shape=(1, number_of_features), maxval=1.0)
    mask = tf.cast(tf.math.greater(probabilities, sample), dtype=tf.float32)
    mask = tf.repeat(mask, number_of_nodes, axis=0)
    anchor_node_features = node_embeddings * (1 - mask)
    sources = tf.range(number_of_nodes, dtype=tf.int32)
    targets = tf.range(number_of_nodes, dtype=tf.int32)
    edge_weights = tf.ones([number_of_nodes], dtype=tf.float32)
    if self._tau > 0.0:
      sources = tf.concat([sources, graph_structure.sources], axis=0)
      targets = tf.concat([targets, graph_structure.targets], axis=0)
      edge_weights = tf.concat(
          [edge_weights, self._tau * graph_structure.weights], axis=0
      )
    return self._graph_data.as_graph_tensor_given_adjacency(
        tf.concat(
            [tf.expand_dims(sources, 0), tf.expand_dims(targets, 0)], axis=0
        ),
        edge_weights=edge_weights,
        node_features=anchor_node_features,
    )
//  [python]
def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
    """
    Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is
    True` else after `sep_token_id`.
    """
    assert shape_list(sep_token_indices)[1] == 2, "`input_ids` should have two dimensions"
    question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1][:, None]
    # bool attention mask with True in locations of global attention
    attention_mask = tf.expand_dims(tf.range(input_ids_shape[1], dtype=tf.int64), axis=0)
    attention_mask = tf.tile(attention_mask, (input_ids_shape[0], 1))
    if before_sep_token is True:
        question_end_index = tf.tile(question_end_index, (1, input_ids_shape[1]))
        attention_mask = tf.to_int64(attention_mask < question_end_index)
    else:
        # last token is separation token and should not be counted and in the middle are two separation tokens
        question_end_index = tf.tile(question_end_index + 1, (1, input_ids_shape[1]))
        attention_mask = tf.to_int64(
            attention_mask > question_end_index,
        ) * tf.to_int64(attention_mask < input_ids_shape[-1])

    return attention_mask
//  [python]
def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
    """
    Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is
    True` else after `sep_token_id`.
    """
    assert shape_list(sep_token_indices)[1] == 2, "`input_ids` should have two dimensions"
    question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1][:, None]
    # bool attention mask with True in locations of global attention
    attention_mask = tf.expand_dims(tf.range(input_ids_shape[1], dtype=tf.int64), axis=0)
    attention_mask = tf.tile(attention_mask, (input_ids_shape[0], 1))
    if before_sep_token is True:
        question_end_index = tf.tile(question_end_index, (1, input_ids_shape[1]))
        attention_mask = tf.cast(attention_mask < question_end_index, dtype=question_end_index.dtype)
    else:
        # last token is separation token and should not be counted and in the middle are two separation tokens
        question_end_index = tf.tile(question_end_index + 1, (1, input_ids_shape[1]))
        attention_mask = tf.cast(
            attention_mask > question_end_index,
            dtype=question_end_index.dtype,
        ) * tf.cast(attention_mask < input_ids_shape[-1], dtype=question_end_index.dtype)

    return attention_mask
//  [python]
def auxiliary_loss(h_states, click_seq, noclick_seq, mask, stag=None):
    #:param h_states:
    #:param click_seq:
    #:param noclick_seq: #[B,T-1,E]
    #:param mask:#[B,1]
    #:param stag:
    #:return:
    hist_len, _ = click_seq.get_shape().as_list()[1:]
    mask = tf.sequence_mask(mask, hist_len)
    mask = mask[:, 0, :]

    mask = tf.to_float(mask)

    click_input_ = tf.concat([h_states, click_seq], -1)

    noclick_input_ = tf.concat([h_states, noclick_seq], -1)

    auxiliary_nn = DNN([100, 50, 1], activation='sigmoid')

    click_prop_ = auxiliary_nn(click_input_, stag=stag)[:, :, 0]

    noclick_prop_ = auxiliary_nn(noclick_input_, stag=stag)[
                    :, :, 0]  # [B,T-1]

    try:
        click_loss_ = - tf.reshape(tf.log(click_prop_),
                                   [-1, tf.shape(click_seq)[1]]) * mask
    except AttributeError:
        click_loss_ = - tf.reshape(tf.compat.v1.log(click_prop_),
                                   [-1, tf.shape(click_seq)[1]]) * mask
    try:
        noclick_loss_ = - \
                            tf.reshape(tf.log(1.0 - noclick_prop_),
                                       [-1, tf.shape(noclick_seq)[1]]) * mask
    except AttributeError:
        noclick_loss_ = - \
                            tf.reshape(tf.compat.v1.log(1.0 - noclick_prop_),
                                       [-1, tf.shape(noclick_seq)[1]]) * mask

    loss_ = reduce_mean(click_loss_ + noclick_loss_)

    return loss_
//  [python]
def auxiliary_loss(h_states, click_seq, noclick_seq, mask, stag=None):
    #:param h_states:
    #:param click_seq:
    #:param noclick_seq: #[B,T-1,E]
    #:param mask:#[B,1]
    #:param stag:
    #:return:
    hist_len, _ = click_seq.get_shape().as_list()[1:]
    mask = tf.sequence_mask(mask, hist_len)
    mask = mask[:, 0, :]

    mask = tf.cast(mask, tf.float32)

    click_input_ = tf.concat([h_states, click_seq], -1)

    noclick_input_ = tf.concat([h_states, noclick_seq], -1)

    auxiliary_nn = DNN([100, 50, 1], activation='sigmoid')

    click_prop_ = auxiliary_nn(click_input_, stag=stag)[:, :, 0]

    noclick_prop_ = auxiliary_nn(noclick_input_, stag=stag)[
                    :, :, 0]  # [B,T-1]

    try:
        click_loss_ = - tf.reshape(tf.log(click_prop_),
                                   [-1, tf.shape(click_seq)[1]]) * mask
    except AttributeError:
        click_loss_ = - tf.reshape(tf.compat.v1.log(click_prop_),
                                   [-1, tf.shape(click_seq)[1]]) * mask
    try:
        noclick_loss_ = - \
                            tf.reshape(tf.log(1.0 - noclick_prop_),
                                       [-1, tf.shape(noclick_seq)[1]]) * mask
    except AttributeError:
        noclick_loss_ = - \
                            tf.reshape(tf.compat.v1.log(1.0 - noclick_prop_),
                                       [-1, tf.shape(noclick_seq)[1]]) * mask

    loss_ = reduce_mean(click_loss_ + noclick_loss_)

    return loss_
//  [python]
def _get_global_attn_indices(is_index_global_attn):
        """compute global attn indices required throughout forward pass"""
        # helper variable
        num_global_attn_indices = tf.math.count_nonzero(is_index_global_attn, axis=1)
        num_global_attn_indices = tf.cast(num_global_attn_indices, dtype=tf.constant(1).dtype)

        # max number of global attn indices in batch
        max_num_global_attn_indices = tf.reduce_max(num_global_attn_indices)

        # indices of global attn
        is_index_global_attn_nonzero = tf.where(is_index_global_attn)

        # helper variable
        is_local_index_global_attn = tf.range(max_num_global_attn_indices) < tf.expand_dims(
            num_global_attn_indices, axis=-1
        )

        # location of the non-padding values within global attention indices
        is_local_index_global_attn_nonzero = tf.where(is_local_index_global_attn)

        # location of the padding values within global attention indices
        is_local_index_no_global_attn_nonzero = tf.where(tf.math.logical_not(is_local_index_global_attn))

        return (
            max_num_global_attn_indices,
            is_index_global_attn_nonzero,
            is_local_index_global_attn_nonzero,
            is_local_index_no_global_attn_nonzero,
        )
//  [python]
def _create_autosummary_var(name, value_expr):
    assert not _autosummary_finalized
    v = t.cast(value_expr, t.float32)
    if v.shape.ndims is 0:
        v = [v, np.float32(1.0)]
    elif v.shape.ndims is 1:
        v = [tf.reduce_sum(v), t.cast(tf.shape(v)[0], t.float32)]
    else:
        v = [tf.reduce_sum(v), tf.reduce_prod(t.cast(tf.shape(v), t.float32))]
    v = tf.cond(tf.is_finite(v[0]), lambda: tf.stack(v), lambda: tf.zeros(2))
    with tf.control_dependencies(None):
        var = tf.Variable(tf.zeros(2)) # [numerator, denominator]
    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
    if name in _autosummary_vars:
        _autosummary_vars[name].append(var)
    else:
        _autosummary_vars[name] = [var]
    return update_op
//  [python]
def _create_autosummary_var(name, value_expr):
    assert not _autosummary_finalized
    v = tf.cast(value_expr, tf.float32)
    if v.shape.ndims is 0:
        v = [v, np.float32(1.0)]
    elif v.shape.ndims is 1:
        v = [tf.reduce_sum(v), tf.cast(tf.shape(v)[0], tf.float32)]
    else:
        v = [tf.reduce_sum(v), tf.reduce_prod(tf.cast(tf.shape(v), tf.float32))]
    v = tf.cond(tf.is_finite(v[0]), lambda: tf.stack(v), lambda: tf.zeros(2))
    with tf.control_dependencies(None):
        var = tf.Variable(tf.zeros(2)) # [numerator, denominator]
    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
    if name in _autosummary_vars:
        _autosummary_vars[name].append(var)
    else:
        _autosummary_vars[name] = [var]
    return update_op
//  [python]
def _concat_with_global_key_attn_probs(
        self,
        attn_scores,
        key_vectors,
        query_vectors,
        max_num_global_attn_indices,
        is_index_global_attn_nonzero,
        is_local_index_global_attn_nonzero,
        is_local_index_no_global_attn_nonzero,
    ):
        batch_size = shape_list(key_vectors)[0]

        # select global key vectors
        global_key_vectors = tf.gather_nd(key_vectors, is_index_global_attn_nonzero)

        # create only global key vectors
        key_vectors_only_global = tf.scatter_nd(
            is_local_index_global_attn_nonzero,
            global_key_vectors,
            shape=(
                batch_size,
                max_num_global_attn_indices,
                self.num_heads,
                self.head_dim,
            ),
        )

        # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
        attn_probs_from_global_key = tf.einsum("blhd,bshd->blhs", query_vectors, key_vectors_only_global)

        # (batch_size, max_num_global_attn_indices, seq_len, num_heads)
        attn_probs_from_global_key_trans = tf.transpose(attn_probs_from_global_key, (0, 3, 1, 2))
        mask_shape = (shape_list(is_local_index_no_global_attn_nonzero)[0],) + tuple(
            shape_list(attn_probs_from_global_key_trans)[-2:]
        )
        mask = tf.ones(mask_shape) * -10000.0
        mask = tf.to_float(mask)

        # scatter mask
        attn_probs_from_global_key_trans = tf.tensor_scatter_nd_update(
            attn_probs_from_global_key_trans,
            is_local_index_no_global_attn_nonzero,
            mask,
        )

        # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
        attn_probs_from_global_key = tf.transpose(attn_probs_from_global_key_trans, (0, 2, 3, 1))

        # concat to attn_probs
        # (batch_size, seq_len, num_heads, extra attention count + 2*window+1)
        attn_scores = tf.concat((attn_probs_from_global_key, attn_scores), axis=-1)

        return attn_scores
//  [python]
def _concat_with_global_key_attn_probs(
        self,
        attn_scores,
        key_vectors,
        query_vectors,
        max_num_global_attn_indices,
        is_index_global_attn_nonzero,
        is_local_index_global_attn_nonzero,
        is_local_index_no_global_attn_nonzero,
    ):
        batch_size = shape_list(key_vectors)[0]

        # select global key vectors
        global_key_vectors = tf.gather_nd(key_vectors, is_index_global_attn_nonzero)

        # create only global key vectors
        key_vectors_only_global = tf.scatter_nd(
            is_local_index_global_attn_nonzero,
            global_key_vectors,
            shape=(
                batch_size,
                max_num_global_attn_indices,
                self.num_heads,
                self.head_dim,
            ),
        )

        # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
        attn_probs_from_global_key = tf.einsum("blhd,bshd->blhs", query_vectors, key_vectors_only_global)

        # (batch_size, max_num_global_attn_indices, seq_len, num_heads)
        attn_probs_from_global_key_trans = tf.transpose(attn_probs_from_global_key, (0, 3, 1, 2))
        mask_shape = (shape_list(is_local_index_no_global_attn_nonzero)[0],) + tuple(
            shape_list(attn_probs_from_global_key_trans)[-2:]
        )
        mask = tf.ones(mask_shape) * -10000.0
        mask = tf.cast(mask, dtype=attn_probs_from_global_key_trans.dtype)

        # scatter mask
        attn_probs_from_global_key_trans = tf.tensor_scatter_nd_update(
            attn_probs_from_global_key_trans,
            is_local_index_no_global_attn_nonzero,
            mask,
        )

        # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
        attn_probs_from_global_key = tf.transpose(attn_probs_from_global_key_trans, (0, 2, 3, 1))

        # concat to attn_probs
        # (batch_size, seq_len, num_heads, extra attention count + 2*window+1)
        attn_scores = tf.concat((attn_probs_from_global_key, attn_scores), axis=-1)

        return attn_scores
//  [python]
def _pad_to_window_size(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        position_ids,
        inputs_embeds,
        pad_token_id,
    ):
        """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
        # padding
        attention_window = (
            self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
        )

        assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"

        input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
        batch_size, seq_len = input_shape[:2]
        padding_len = (attention_window - seq_len % attention_window) % attention_window

        if padding_len > 0:
            logger.info(
                f"Input ids are automatically padded from {seq_len} to {seq_len + padding_len} to be a multiple of "
                f"`config.attention_window`: {attention_window}"
            )

        paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])

        if input_ids is not None:
            input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)

        if position_ids is not None:
            # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
            position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)

        if inputs_embeds is not None:

            def pad_embeddings():
                input_ids_padding = tf.fill((batch_size, padding_len), self.pad_token_id)
                inputs_embeds_padding = self.embeddings(input_ids_padding)
                return tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)

            inputs_embeds = tf.cond(tf.math.greater(padding_len, 0), pad_embeddings, lambda: inputs_embeds)

        attention_mask = tf.pad(attention_mask, paddings, constant_values=False)  # no attention on the padding tokens
        token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0)  # pad with token_type_id = 0

        return (
            padding_len,
            input_ids,
            attention_mask,
            token_type_ids,
            position_ids,
            inputs_embeds,
        )
//  [python]
def _pad_to_window_size(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        position_ids,
        inputs_embeds,
        pad_token_id,
    ):
        """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
        # padding
        attention_window = (
            self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
        )

        assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"

        input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
        batch_size, seq_len = input_shape[:2]
        padding_len = (attention_window - seq_len % attention_window) % attention_window

        if padding_len > 0:
            logger.info(
                f"Input ids are automatically padded from {seq_len} to {seq_len + padding_len} to be a multiple of "
                f"`config.attention_window`: {attention_window}"
            )

        paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])

        if input_ids is not None:
            input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)

        if position_ids is not None:
            # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
            position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)

        if inputs_embeds is not None:

            def pad_embeddings():
                input_ids_padding = tf.cast(tf.fill((batch_size, padding_len), self.pad_token_id), tf.int64)
                inputs_embeds_padding = self.embeddings(input_ids_padding)
                return tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)

            inputs_embeds = tf.cond(tf.math.greater(padding_len, 0), pad_embeddings, lambda: inputs_embeds)

        attention_mask = tf.pad(attention_mask, paddings, constant_values=False)  # no attention on the padding tokens
        token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0)  # pad with token_type_id = 0

        return (
            padding_len,
            input_ids,
            attention_mask,
            token_type_ids,
            position_ids,
            inputs_embeds,
        )
//  [python]
def register_gradients(self, loss, vars):
        assert not self._updates_applied

        # Validate arguments.
        if isinstance(vars, dict):
            vars = list(vars.values()) # allow passing in Network.trainables as vars
        assert isinstance(vars, list) and len(vars) >= 1
        assert all(is_tf_expression(expr) for expr in vars + [loss])
        if self._grad_shapes is None:
            self._grad_shapes = [shape_to_list(var.shape) for var in vars]
        assert len(vars) == len(self._grad_shapes)
        assert all(shape_to_list(var.shape) == var_shape for var, var_shape in zip(vars, self._grad_shapes))
        dev = loss.device
        assert all(var.device == dev for var in vars)

        # Register device and compute gradients.
        with tf.name_scope(self.id + '_grad'), tf.device(dev):
            if dev not in self._dev_opt:
                opt_name = self.scope.replace('/', '_') + '_opt%d' % len(self._dev_opt)
                self._dev_opt[dev] = self.optimizer_class(name=opt_name, learning_rate=self.learning_rate, **self.optimizer_kwargs)
                self._dev_grads[dev] = []
            loss = self.apply_loss_scaling(tf.to_float(loss))
            grads = self._dev_opt[dev].compute_gradients(loss, vars, gate_gradients=tf.train.Optimizer.GATE_NONE) # disable gating to reduce memory usage
            grads = [(g, v) if g is not None else (tf.zeros_like(v), v) for g, v in grads] # replace disconnected gradients with zeros
            self._dev_grads[dev].append(grads)
//  [python]
def register_gradients(self, loss, vars):
        assert not self._updates_applied

        # Validate arguments.
        if isinstance(vars, dict):
            vars = list(vars.values()) # allow passing in Network.trainables as vars
        assert isinstance(vars, list) and len(vars) >= 1
        assert all(is_tf_expression(expr) for expr in vars + [loss])
        if self._grad_shapes is None:
            self._grad_shapes = [shape_to_list(var.shape) for var in vars]
        assert len(vars) == len(self._grad_shapes)
        assert all(shape_to_list(var.shape) == var_shape for var, var_shape in zip(vars, self._grad_shapes))
        dev = loss.device
        assert all(var.device == dev for var in vars)

        # Register device and compute gradients.
        with tf.name_scope(self.id + '_grad'), tf.device(dev):
            if dev not in self._dev_opt:
                opt_name = self.scope.replace('/', '_') + '_opt%d' % len(self._dev_opt)
                self._dev_opt[dev] = self.optimizer_class(name=opt_name, learning_rate=self.learning_rate, **self.optimizer_kwargs)
                self._dev_grads[dev] = []
            loss = self.apply_loss_scaling(tf.cast(loss, tf.float32))
            grads = self._dev_opt[dev].compute_gradients(loss, vars, gate_gradients=tf.train.Optimizer.GATE_NONE) # disable gating to reduce memory usage
            grads = [(g, v) if g is not None else (tf.zeros_like(v), v) for g, v in grads] # replace disconnected gradients with zeros
            self._dev_grads[dev].append(grads)
//  [python]
def encode(self, x, encode_params):
    """See base class."""
    clipped_x = tf.clip_by_value(
        x,
        t.cast(encode_params[self.MIN_PARAMS_KEY], x.dtype),
        t.cast(encode_params[self.MAX_PARAMS_KEY], x.dtype))
    return {self.ENCODED_VALUES_KEY: clipped_x}
//  [python]
def encode(self, x, encode_params):
    """See base class."""
    clipped_x = tf.clip_by_value(
        x,
        tf.cast(encode_params[self.MIN_PARAMS_KEY], x.dtype),
        tf.cast(encode_params[self.MAX_PARAMS_KEY], x.dtype))
    return {self.ENCODED_VALUES_KEY: clipped_x}
//  [python]
def sample(self, amount=None):
    amount = [amount] if amount else []
    indices = self._dist.sample(*amount)
    sample = self._one_hot(indices)
    probs = self._dist.probs_parameter()
    sample += t.cast(probs - tf.stop_gradient(probs), self._dtype)
    return sample
//  [python]
def sample(self, amount=None):
    amount = [amount] if amount else []
    indices = self._dist.sample(*amount)
    sample = self._one_hot(indices)
    probs = self._dist.probs_parameter()
    sample += tf.cast(probs - tf.stop_gradient(probs), self._dtype)
    return sample
//  [python]
def _inverse(self, y):
    dtype = y.dtype
    y = t.cast(y, tf.float32)
    y = tf.where(
        tf.less_equal(tf.abs(y), 1.),
        tf.clip_by_value(y, -0.99999997, 0.99999997), y)
    y = tf.atanh(y)
    y = t.cast(y, dtype)
    return y
//  [python]
def _inverse(self, y):
    dtype = y.dtype
    y = tf.cast(y, tf.float32)
    y = tf.where(
        tf.less_equal(tf.abs(y), 1.),
        tf.clip_by_value(y, -0.99999997, 0.99999997), y)
    y = tf.atanh(y)
    y = tf.cast(y, dtype)
    return y
//  [python]
def _attn(self, q, k, v):
        w = tf.matmul(q, k)
        if self.scale:
            w = w / tf.sqrt(tf.to_float(v.shape[-1]))
        # self.b may be larger than w, so we need to crop it
        b = self.b[:, :, :w.shape[-2], :w.shape[-1]]
        w = w * b + 1e-9 * (1 - b)
        w = tf.nn.softmax(w, -1)
        return tf.matmul(w, v)
//  [python]
def _attn(self, q, k, v):
        w = tf.matmul(q, k)
        if self.scale:
            w = w / tf.sqrt(tf.cast(v.shape[-1], tf.float32))
        # self.b may be larger than w, so we need to crop it
        b = self.b[:, :, :w.shape[-2], :w.shape[-1]]
        w = w * b + 1e-9 * (1 - b)
        w = tf.nn.softmax(w, -1)
        return tf.matmul(w, v)
//  [python]
def _decode(self, record: tf.Tensor) -> Dict[str, tf.Tensor]:
    """Decodes a serialized tf.Example."""

    name_to_features = {
        'input_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_mask': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'segment_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'label_ids': tf.io.FixedLenFeature([], tf.int64),
    }
    example = tf.io.parse_single_example(record, name_to_features)

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

    return example
//  [python]
def _decode(self, record: tf.Tensor) -> Dict[str, tf.Tensor]:
    """Decodes a serialized tf.Example."""

    name_to_features = {
        'input_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_mask': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'segment_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'label_ids': tf.io.FixedLenFeature([], tf.int64),
    }
    example = tf.io.parse_single_example(record, name_to_features)

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

    return example
//  [python]
def load_datum(tfrecord_dict):
  """Convert tfrecord dictionary to tensors."""
  pose_body = TensorflowPoseBody.from_tfrecord(tfrecord_dict)
  pose = Pose(header=get_openpose_header(), body=pose_body)
  tgt = tf.io.decode_raw(tfrecord_dict["is_signing"], out_type=tf.int8)

  fps = pose.body.fps
  frames = tf.cast(tf.size(tgt), dtype=fps.dtype)

  # Get only relevant input components
  pose = pose.get_components(FLAGS.input_components)

  return {
      "fps": pose.body.fps,
      "frames": frames,
      "tgt": tgt,
      "pose_data_tensor": pose.body.data.tensor,
      "pose_data_mask": pose.body.data.mask,
      "pose_confidence": pose.body.confidence,
  }
//  [python]
def tensors_to_item(self, keys_to_tensors):
    return t.cast(
        self._num_of_views * keys_to_tensors[self._original_width_key] /
        keys_to_tensors[self._width_key], dtype=t.int64)
//  [python]
def tensors_to_item(self, keys_to_tensors):
    return tf.cast(
        self._num_of_views * keys_to_tensors[self._original_width_key] /
        keys_to_tensors[self._width_key], dtype=tf.int64)
//  [python]
def _pad_to_window_size(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        position_ids,
        inputs_embeds,
        pad_token_id,
    ):
        """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
        # padding
        attention_window = (
            self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
        )

        assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"

        input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
        batch_size, seq_len = input_shape[:2]
        padding_len = (attention_window - seq_len % attention_window) % attention_window

        paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])

        if input_ids is not None:
            input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)

        if position_ids is not None:
            # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
            position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)

        if inputs_embeds is not None:
            if padding_len > 0:
                input_ids_padding = tf.fill((batch_size, padding_len), self.pad_token_id)
                inputs_embeds_padding = self.embeddings(input_ids_padding)
                inputs_embeds = tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)

        attention_mask = tf.pad(attention_mask, paddings, constant_values=False)  # no attention on the padding tokens
        token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0)  # pad with token_type_id = 0

        return (
            padding_len,
            input_ids,
            attention_mask,
            token_type_ids,
            position_ids,
            inputs_embeds,
        )