import torch


def _pdist(a, b):
    """Compute pair-wise squared distance between points in `a` and `b`.

    Parameters
    ----------
    a : array_like
        An NxM matrix of N samples of dimensionality M.
    b : array_like
        An LxM matrix of L samples of dimensionality M.

    Returns
    -------
    ndarray
        Returns a matrix of size len(a), len(b) such that eleement (i, j)
        contains the squared distance between `a[i]` and `b[j]`.

    """
    if len(a) == 0 or len(b) == 0:
        return torch.zeros(len(a), len(b), dtype=a.dtype, device=a.device)

    a = a[:, None, :]
    b = b[None, :, :]

    return torch.sum(torch.pow(a - b, 2), dim=-1)