import torch


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 (:obj:`torch.Tensor` of shape (B1, ..., Bn, num_segments, V1, ...)):
            Tensor with segment values.
        index (:obj:`IndexMap` of shape (B1, ..., Bn, I1, ..., Ik)):
            IndexMap.
        name (:obj:`str`, `optional`, defaults to 'segmented_gather'):
            Name for the operation. Currently not used

    Returns:
        :obj:`tuple(torch.Tensor)`: Tensor of shape (B1, ..., Bn, I1, ..., Ik, V1, ...) with the gathered values.
    """
    indices = index.indices
    # first, check whether the indices of the index represent scalar values (i.e. not vectorized)
    if len(values.shape[index.batch_dims :]) < 2:
        return torch.gather(
            values,
            index.batch_dims,
            indices.view(
                values.size()[0], -1
            ),  # torch.gather expects index to have the same number of dimensions as values
        ).view(indices.size())
    else:
        # this means we have a vectorized version
        # we have to adjust the index
        indices = indices.unsqueeze(-1).expand(values.shape)
        return torch.gather(values, index.batch_dims, indices)