def flatten(index, name="segmented_flatten"):
    """
    Flattens a batched index map to a 1d index map. This operation relabels the segments to keep batch elements
    distinct. The k-th batch element will have indices shifted by `num_segments` * (k - 1). The result is a tensor with
    `num_segments` multiplied by the number of elements in the batch.

    Args:
      index: IndexMap to flatten.
      name: Name for the TensorFlow operation.

    Returns:
      The flattened IndexMap.
    """
    batch_size = tf.reduce_prod(index.batch_shape())
    offset = tf.range(batch_size) * index.num_segments
    offset = tf.reshape(offset, index.batch_shape())
    for _ in range(index.batch_dims, index.indices.shape.rank):
        offset = tf.expand_dims(offset, -1)

    indices = tf.cast(offset, index.indices.dtype) + index.indices
    return IndexMap(indices=tf.reshape(indices, [-1]), num_segments=index.num_segments * batch_size, batch_dims=0)

    Flattens a batched index map to a 1d index map. This operation relabels the segments to keep batch elements
    distinct. The k-th batch element will have indices shifted by `num_segments` * (k - 1). The result is a tensor with
    `num_segments` multiplied by the number of elements in the batch.

    Args:
      index: IndexMap to flatten.
      name: Name for the TensorFlow operation.

    Returns:
      The flattened IndexMap.
    
dpkd/transformers/src/transformers/models/tapas/modeling_tf_tapas.py