def segment_indices(segment_ids: tf.Tensor,
                    name: Optional[str] = None) -> tf.Tensor:
  """Returns a `Tensor` of indices within each segment.

  segment_ids should be a sequence of non-decreasing non-negative integers that
  define a set of segments, e.g. [0, 0, 1, 2, 2, 2] defines 3 segments of length
  2, 1 and 3.  The return value is a `Tensor` containing the indices within each
  segment.

  Example:

  >>> result = tft.segment_indices(tf.constant([0, 0, 1, 2, 2, 2]))
  >>> print(result)
  tf.Tensor([0 1 0 0 1 2], shape=(6,), dtype=int32)

  Args:
    segment_ids: A 1-d `Tensor` containing an non-decreasing sequence of
      non-negative integers with type `tf.int32` or `tf.int64`.
    name: (Optional) A name for this operation.

  Returns:
    A `Tensor` containing the indices within each segment.
  """
  ndims = segment_ids.get_shape().ndims
  if ndims != 1 and ndims is not None:
    raise ValueError(
        'segment_indices requires a 1-dimensional input. '
        'segment_indices has {} dimensions.'.format(ndims))
  with tf.compat.v1.name_scope(name, 'segment_indices'):
    # TODO(KesterTong): This is a fundamental operation for segments, write a C++
    # op to do this.
    # TODO(KesterTong): Add a check that segment_ids are increasing.
    segment_lengths = tf.math.segment_sum(
        tf.ones_like(segment_ids), segment_ids)
    segment_starts = tf.gather(tf.concat([[0], tf.cumsum(segment_lengths)], 0),
                               segment_ids)
    return (tf.range(tf.size(input=segment_ids, out_type=segment_ids.dtype)) -
            segment_starts)
Returns a `Tensor` of indices within each segment.

  segment_ids should be a sequence of non-decreasing non-negative integers that
  define a set of segments, e.g. [0, 0, 1, 2, 2, 2] defines 3 segments of length
  2, 1 and 3.  The return value is a `Tensor` containing the indices within each
  segment.

  Example:

  >>> result = tft.segment_indices(tf.constant([0, 0, 1, 2, 2, 2]))
  >>> print(result)
  tf.Tensor([0 1 0 0 1 2], shape=(6,), dtype=int32)

  Args:
    segment_ids: A 1-d `Tensor` containing an non-decreasing sequence of
      non-negative integers with type `tf.int32` or `tf.int64`.
    name: (Optional) A name for this operation.

  Returns:
    A `Tensor` containing the indices within each segment.
  
tensorflow_transform/mappers.py