import torch


def pdist2sq(X, Y):
    """ Computes the squared Euclidean distance between all pairs x in X, y in Y """
    C = -2 * torch.matmul(X, Y.T)
    nx = torch.sum(X ** 2, dim=1).unsqueeze(1)
    ny = torch.sum(Y ** 2, dim=1).unsqueeze(1)
    D = (C + ny.T) + nx
    return D