def get_dense_span_ends_from_starts(dense_span_starts: tf.Tensor,
                                    dense_span_ends: tf.Tensor) -> tf.Tensor:
  """For every mention start positions finds the corresponding end position."""
  seq_len = tf.shape(dense_span_starts)[0]
  start_pos = tf.cast(tf.where(tf.equal(dense_span_starts, 1)), tf.int32)
  end_pos = tf.cast(
      tf.squeeze(tf.where(tf.equal(dense_span_ends, 1)), 1), tf.int32)
  dense_span_ends_from_starts = tf.zeros(seq_len, dtype=tf.int32)
  dense_span_ends_from_starts = tf.tensor_scatter_nd_add(
      dense_span_ends_from_starts, start_pos, end_pos)
  return dense_span_ends_from_starts
For every mention start positions finds the corresponding end position.
language/mentionmemory/utils/mention_preprocess_utils.py