def _batched_range(batch_size: int, length: int, axis: int,
                   dtype: tf.dtypes.DType) -> tf.Tensor:
  """Produces multiple tf.range stacked along the `axis` `batch_size` times."""
  if axis == 0:
    repeats = [batch_size, 1]
  elif axis == 1:
    repeats = [1, batch_size]
  else:
    raise ValueError('`_batched_range` only accepts axis argument being 0 or 1')
  return tf.tile(tf.expand_dims(tf.range(length, dtype=dtype), axis), repeats)
Produces multiple tf.range stacked along the `axis` `batch_size` times.
language/mentionmemory/utils/mention_preprocess_utils.py