import torch


def _cosine_distance(a, b, norm=True):
    """
    Compute pair-wise cosine distance between points in `a` and `b`.
    """
    # return 1. - np.dot(a, b.T)
    if norm:
        a = torch.nn.functional.normalize(a, dim=1, p=2)
        b = torch.nn.functional.normalize(b, dim=1, p=2)
    return 1. - torch.mm(a, b.t())