def process_batchwise_mention_targets(
    dense_span_starts: tf.Tensor,
    dense_span_ends: tf.Tensor,
    dense_mention_ids: tf.Tensor,
    dense_linked_mention_mask: tf.Tensor,
    dense_is_masked: tf.Tensor,
    max_mentions: int,
    max_mention_targets: int,
) -> Dict[str, tf.Tensor]:
  """Processes mention targets and subsamples/pads as necessary.

  This function does two things. First, it selects which mentions to mark as
  mentions for mention-aware text encoders (in case the number of mentions
  exceeds the max number of mentions). Second, it selects which linked
  mentions to use as targets for mention objectives. To reduce subsampling and
  padding, the function operates over all mentions in a batch, generating
  flattened arrays. The encoder reconstructs the original mention positions
  from an array which specifies each mention's position in the batch. Linked
  mentions are given priority for sampling.

  Args:
    dense_span_starts: dense mention start positions.
    dense_span_ends: dense mention end positions.
    dense_mention_ids: dense entity ids for linked mentions in passage.
    dense_linked_mention_mask: dense mask for linked mentions in passage.
    dense_is_masked: dense mask for masked positions in passage.
    max_mentions: max number of mentions to be considered in model.
    max_mention_targets: max number of mentions to be used for linking loss.

  Returns:
    Mention starts, mention ends, mention mask,
    mention target indices (into start/end positions),
    mention target ids, mention target weights, mention_target_batch_positions,
    mention_target_start_positions, mention_target_end_positions
  """

  seq_len = tf.shape(dense_span_starts)[1]

  # The linking mask has 1's for every part of the mention, we only
  # want it for starts...
  linking_mask_start_indexed = dense_span_starts * dense_linked_mention_mask

  # values in {0, 1, 2}:
  # 0: not a masking location.
  # 1: a masking location.
  # 2: a masking and linking location.
  prioritized_span_starts = dense_span_starts + linking_mask_start_indexed
  prioritized_span_starts = tf.cast(prioritized_span_starts, tf.float32)

  # Add random [0; 1) values for a uniform sampling in case
  # there are more mention than `max_mentions`
  prioritized_span_starts += tf.random.uniform(
      tf.shape(prioritized_span_starts))

  _, global_start_indices = tf.math.top_k(
      _flatten(prioritized_span_starts), k=max_mentions)

  dense_span_starts_flatten = _flatten(dense_span_starts)
  dense_span_ends_at_starts = get_dense_span_ends_from_starts(
      dense_span_starts_flatten, _flatten(dense_span_ends))
  global_end_indices = tf.gather(dense_span_ends_at_starts,
                                 global_start_indices)

  dtype = dense_span_starts.dtype
  mention_batch_positions = tf.math.floordiv(global_start_indices, seq_len)
  mention_batch_positions = tf.cast(mention_batch_positions, dtype=dtype)
  mention_start_positions = tf.math.floormod(global_start_indices, seq_len)
  mention_start_positions = tf.cast(mention_start_positions, dtype=dtype)
  mention_end_positions = tf.math.floormod(global_end_indices, seq_len)
  mention_end_positions = tf.cast(mention_end_positions, dtype=dtype)
  mention_mask = tf.gather(dense_span_starts_flatten, global_start_indices)
  mention_mask = tf.cast(mention_mask, dtype=dtype)
  mention_batch_positions *= mention_mask
  mention_start_positions *= mention_mask
  mention_end_positions *= mention_mask

  mention_target_weights = tf.gather(
      _flatten(linking_mask_start_indexed), global_start_indices)
  mention_target_weights = mention_target_weights[:max_mention_targets]
  mention_target_weights = tf.cast(mention_target_weights, dtype=dtype)
  mention_target_indices = tf.range(max_mention_targets, dtype=dtype)
  mention_target_indices = mention_target_indices * mention_target_weights
  mention_target_ids = tf.gather(
      _flatten(dense_mention_ids), global_start_indices)
  mention_target_ids = mention_target_ids[:max_mention_targets]
  mention_target_ids = tf.cast(mention_target_ids, dtype=dtype)
  mention_target_ids = mention_target_ids * mention_target_weights
  indices = tf.stack((mention_batch_positions, mention_start_positions), axis=1)
  mention_is_masked = tf.gather_nd(dense_is_masked, indices)
  mention_target_is_masked = mention_is_masked[:max_mention_targets]

  features = {
      'mention_batch_positions': mention_batch_positions,
      'mention_start_positions': mention_start_positions,
      'mention_end_positions': mention_end_positions,
      'mention_mask': mention_mask,
      'mention_is_masked': mention_is_masked,
      'mention_target_ids': mention_target_ids,
      'mention_target_indices': mention_target_indices,
      'mention_target_is_masked': mention_target_is_masked,
  }
  mention_target_features = prepare_mention_target_features(
      mention_batch_positions, mention_start_positions, mention_end_positions,
      mention_mask, mention_target_weights, mention_target_indices)
  features.update(mention_target_features)
  return features
Processes mention targets and subsamples/pads as necessary.

  This function does two things. First, it selects which mentions to mark as
  mentions for mention-aware text encoders (in case the number of mentions
  exceeds the max number of mentions). Second, it selects which linked
  mentions to use as targets for mention objectives. To reduce subsampling and
  padding, the function operates over all mentions in a batch, generating
  flattened arrays. The encoder reconstructs the original mention positions
  from an array which specifies each mention's position in the batch. Linked
  mentions are given priority for sampling.

  Args:
    dense_span_starts: dense mention start positions.
    dense_span_ends: dense mention end positions.
    dense_mention_ids: dense entity ids for linked mentions in passage.
    dense_linked_mention_mask: dense mask for linked mentions in passage.
    dense_is_masked: dense mask for masked positions in passage.
    max_mentions: max number of mentions to be considered in model.
    max_mention_targets: max number of mentions to be used for linking loss.

  Returns:
    Mention starts, mention ends, mention mask,
    mention target indices (into start/end positions),
    mention target ids, mention target weights, mention_target_batch_positions,
    mention_target_start_positions, mention_target_end_positions
  
language/mentionmemory/utils/mention_preprocess_utils.py