import torch


def _squared_distance(x1, x2):
    """
    Compute squared distance matrix.

    :param x1: [N1 x D]
    :type x1: torch.Tensor
    :param x2: [N2 x D]
    :type x2: torch.Tensor

    :return: [N1 x N2] squared distance matrix
    """

    r2 = (
        torch.sum(x1 ** 2, dim=1, keepdim=True)
        - 2.0 * x1 @ x2.t()
        + torch.sum(x2 ** 2, dim=1, keepdim=True).t()
    )
    r2 = r2 - (torch.clamp(r2, max=0.0)).detach()

    return r2