def _get_2d_index(x: tf.Tensor, y: tf.Tensor) -> tf.Tensor:
  """Generates 2D index given separate arrays for each coordinate."""
  # Index is always int32
  return tf.cast(tf.stack([x, y], axis=1), tf.int32)
Generates 2D index given separate arrays for each coordinate.
language/mentionmemory/utils/mention_preprocess_utils.py