src/transformers/models/tapas/modeling_tf_tapas.py/gather
def gather(values, index, name="segmented_gather"):
    """
    Gathers from `values` using the index map. For each element in the domain of the index map this operation looks up
    a value for that index in `values`. Two elements from the same segment always get assigned the same value.

    Args:
      values: [B1, ..., Bn, num_segments, V1, ...] Tensor with segment values.
      index: [B1, ..., Bn, I1, ..., Ik] IndexMap.
      name: Name for the TensorFlow operation.

    Returns:
      [B1, ..., Bn, I1, ..., Ik, V1, ...] Tensor with the gathered values.
    """
    return tf.gather(values, index.indices, batch_dims=index.batch_dims, name=name)