def get_dense_is_inside_for_dense_spans(
    dense_start_positions: tf.Tensor,
    dense_end_positions: tf.Tensor) -> tf.Tensor:
  """Dense mask whether position is inside span given dense starts / ends."""
  # `tf.cumsum(dense_start_positions)[i]` computes how many spans start before
  # or on the i-th position.
  # `tf.cumsum(dense_end_positions, exclusive=True`) computes how many spans
  #  ends strictly before i-th positions.
  # Their difference is equal to how many spans start before i-th position and
  # end after or on i-th position. This is precisely how many spans contain
  # i-th position.
  is_inside_span = (
      tf.cumsum(dense_start_positions) -
      tf.cumsum(dense_end_positions, exclusive=True))
  # Adjust for the case of overlapping spans
  is_inside_span = tf.minimum(is_inside_span, 1)
  return is_inside_span
Dense mask whether position is inside span given dense starts / ends.
language/mentionmemory/utils/mention_preprocess_utils.py