def __init__(self, indices, num_segments, batch_dims=0):
        """
        Creates an index.

        Args:
          indices: <int32> Tensor of indices, same shape as `values`.
          num_segments: <int32> Scalar tensor, the number of segments. All elements
            in a batched segmented tensor must have the same number of segments (although many segments can be empty).
          batch_dims: Python integer, the number of batch dimensions. The first
            `batch_dims` dimensions of a SegmentedTensor are treated as batch dimensions. Segments in different batch
            elements are always distinct even if they have the same index.
        """
        self.indices = tf.convert_to_tensor(indices)
        self.num_segments = tf.convert_to_tensor(num_segments)
        self.batch_dims = batch_dims

        Creates an index.

        Args:
          indices: <int32> Tensor of indices, same shape as `values`.
          num_segments: <int32> Scalar tensor, the number of segments. All elements
            in a batched segmented tensor must have the same number of segments (although many segments can be empty).
          batch_dims: Python integer, the number of batch dimensions. The first
            `batch_dims` dimensions of a SegmentedTensor are treated as batch dimensions. Segments in different batch
            elements are always distinct even if they have the same index.
        
dpkd/transformers/src/transformers/models/tapas/modeling_tf_tapas.py