import torch


def gather(data, index_2d):
    """
    Gather data alongs the second dim. Assume data's shape as (batch_size, dim1, dim2, ...),
    and index_2d's shape is (batch_size, dim1).
    output[i][j] = data[i][index_2d[i][j]]

    This function does not require data, output, or index_2d having the same shape, which
     is mandated by torch.gather.
    """
    batch_size = data.shape[0]
    data_shape = data.shape[2:]
    index_len = index_2d.shape[1]
    res = data[
        torch.arange(batch_size, device=data.device).repeat_interleave(index_len),
        index_2d.flatten(),
    ].view(batch_size, index_len, *data_shape)
    return res