def get_diff_sq(a, b):
    """
    Computes squared pairwise differences between a and b.

    :param a: tensor a
    :param b: tensor b
    :return: squared pairwise differences between a and b
    """
    aa = a.matmul(a.t())
    bb = b.matmul(b.t())
    ab = a.matmul(b.t())
    diff_sq = -2 * ab + aa.diag().unsqueeze(1) + bb.diag().unsqueeze(0)
    return diff_sq