def compute_which_mentions_fit_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,
    new_length: int,
) -> tf.Tensor:
  """Computes a mask for which mentions will fit after adding entity tokens."""

  positions = compute_positions_shift_with_entity_tokens(
      mention_mask, mention_batch_positions, mention_start_positions,
      mention_end_positions, batch_size, old_length)

  def get_new_positions(old_positions):
    index_2d = _get_2d_index(mention_batch_positions, old_positions)
    return tf.gather_nd(positions, index_2d)

  # `get_new_positions(mention_end_positions)` returns new positions for the
  # last token per mention. However, we want `mention_end_positions` to point
  # to the entity end token, so we add 1.
  new_mention_end_positions = get_new_positions(mention_end_positions) + 1

  new_mention_mask = tf.less(new_mention_end_positions, new_length)
  return tf.cast(new_mention_mask, mention_mask.dtype)
Computes a mask for which mentions will fit after adding entity tokens.
language/mentionmemory/utils/mention_preprocess_utils.py