import torch


def batch_gather(input, index):
    """Similar to tf.compat.v1.batch_gather.

    Args:
        input: (Batch, Dim1, Dim2)
        index: (Batch, N_ind)

    Returns:
        output: (Batch, N_ind, Dim2)
            output[i, j, k] = input[i, index[i, j], k]
    """
    index_exp = index.unsqueeze(2).expand(-1, -1, input.size(2))
    out = torch.gather(input, 1, index_exp)  # (Batch, N_ind, Dim2)
    return out