import torch


def knn(dist_mat, k=20):
    """
    Input: pairwise distance: (B, N, N)
           k: int
    Output: nearest neighbors: (B, N, k)
    """

    _, nn_idx = torch.topk(dist_mat, k=k, largest=False, sorted=False)
    return nn_idx