import torch


def cos_dist_mat(x, y):
    """
    Cos similarity
    :param x: [n1, d]
    :param y: [n2, d]
    :return: [n1, n2]
    """
    n1, d1 = x.shape
    n2, d2 = y.shape
    assert d1 == d2

    if n1 * n2 > 10000:
        D = torch.cat(
            [torch.cosine_similarity(x[i].view(1, 1, d1), y.view(1, n2, d2), -1) for i in range(n1)], 0
        )
    else:
        D = torch.cosine_similarity(x.reshape(n1, 1, d1), y.reshape(1, n2, d2), -1)
    return 1 - D