import torch


def pdist(X,Y):
    """ Computes the squared Euclidean distance between all pairs x in X, y in Y """
    # From https://github.com/clinicalml/cfrnet/blob/master/cfr/util.py
    M = -2*torch.matmul(X,torch.transpose(Y,0,1))
    sqx = torch.sum(X**2,1,keepdim=True)
    sqy = torch.sum(Y**2,1,keepdim=True)
    D = M + torch.transpose(sqy,0,1) + sqx
    return D