def gaussian_stein_kernel_single(x, score_x, sigma, return_kernel=False):
    """Compute the Gaussian Stein kernel between x and x


    Parameters
    ----------
    x : torch.tensor, shape (n, p)
        Input particles
    score_x : torch.tensor, shape (n, p)
        The score of x
    sigma : float
        Bandwidth
    return_kernel : bool
        whether the original kernel k(xi, xj) should also be returned

    Return
    ------
    stein_kernel : torch.tensor, shape (n, n)
        The linear Stein kernel
    kernel : torch.tensor, shape (n, n)
        The base kernel, only returned id return_kernel is True
    """
    _, p = x.shape
    # Gaussian kernel:
    norms = (x ** 2).sum(-1)
    dists = -2 * x @ x.t() + norms[:, None] + norms[None, :]
    k = (-dists / 2 / sigma).exp()

    # Dot products:
    diffs = (x * score_x).sum(-1, keepdim=True) - (x @ score_x.t())
    diffs = diffs + diffs.t()
    scalars = score_x.mm(score_x.t())
    der2 = p - dists / sigma
    stein_kernel = k * (scalars + diffs / sigma + der2 / sigma)
    if return_kernel:
        return stein_kernel, k
    return stein_kernel