import torch


def index_points(points, idx):
    """Returns subsamples of points tensor index by idx tensor.

    Args:
        points (torch.Tensor): Shape (batch_size, num_points, feature_dim)
        idx (torch.Tensor): Index tensor of shape (batch_size, num_output_points) OR
            (batch_size, num_output_points, num_neighbors).

    Returns:
        out_points (torch.Tensor): Indexed points, with shape (batch_size, num_output_points, feature_dim)
            OR (batch_size, num_output_points, num_neighbors, feature_dim).
    """
    batch_size = points.shape[0]
    # Create tensor of batch indices with shape (batch_size,)
    batch_indices = torch.arange(batch_size, dtype=torch.long).to(points.device)
    if len(idx.shape) == 2:
        # Repeat batch indices to match shape of idx tensor
        batch_indices = batch_indices.view((batch_size, 1)).repeat((1,) + idx.shape[1:])
    elif len(idx.shape) == 3:
        # Repeat batch indices to match shape of idx tensor
        batch_indices = batch_indices.view((batch_size, 1, 1)).repeat((1,) + idx.shape[1:])
    # Pytorch tensor indexing magic
    out_points = points[batch_indices, idx, :]
    return out_points