def compute_positions_shift_with_entity_tokens(
    mention_mask: tf.Tensor,
    mention_batch_positions: tf.Tensor,
    mention_start_positions: tf.Tensor,
    mention_end_positions: tf.Tensor,
    batch_size: int,
    old_length: int,
) -> tf.Tensor:
  """Computes the new position for every position in the old sequence."""
  old_shape = (batch_size, old_length)

  def get_positions_shift(positions: tf.Tensor, exclusive: bool) -> tf.Tensor:
    index_2d = _get_2d_index(mention_batch_positions, positions)
    return tf.cumsum(
        tf.scatter_nd(index_2d, mention_mask, old_shape),
        axis=1,
        exclusive=exclusive)

  # We start with identity mapping from old positions to new positions.
  positions = _batched_range(batch_size, old_length, 0, mention_mask.dtype)
  # We need to insert entity tokens BEFORE the first token of every mention.
  # (note `exclusive=False`). Therefore, all the tokens after will need to be
  # shifted right by 1.
  positions += get_positions_shift(mention_start_positions, exclusive=False)
  # We need to insert entity tokens AFTER the last token of every mention.
  # (note `exclusive=True`). Therefore, all the tokens after will need to be
  # shifted right by 1.
  positions += get_positions_shift(mention_end_positions, exclusive=True)
  return positions
Computes the new position for every position in the old sequence.
language/mentionmemory/utils/mention_preprocess_utils.py