import torch


def batched_index_select(inputs, idx):
    """
    This can be used for neighbors features fetching
    faster bached index select, return a feature by a tensor idx.
    :param inputs: torch.Size([batch_size, num_dims, num_vertices, 1])
    :param idx: torch.Size([batch_size, num_vertices, k])
    :return: torch.Size([batch_size, num_dims, num_vertices, k])
    """

    batch_size, num_dims, num_vertices = inputs.shape[:3]
    k = idx.shape[-1]
    idx_base = torch.arange(0, batch_size, device=idx.device).view(-1, 1, 1) * num_vertices
    idx = idx + idx_base
    idx = idx.contiguous().view(-1)

    x = inputs.transpose(2, 1).contiguous()
    feature = x.view(batch_size * num_vertices, -1)[idx, :]
    feature = feature.view(batch_size, num_vertices, k, num_dims).permute(0, 2, 3, 1).contiguous()
    return feature