import torch


def svgd_direction(x, scores_x, sigma):
    """Compute the SVGD direction at x with a Gaussian kernel


    Parameters
    ----------
    x : torch.tensor, shape (n, p)
        Input particles
    score_x : torch.tensor, shape (n, p)
        The score of x
    sigma : float
        Bandwidth

    Return
    ------
    direction : torch.tensor, shape (n, p)
        the SVDG direction
    """
    _, p = x.shape
    d = x[:, None, :] - x[None, :, :]
    dists = (d ** 2).sum(axis=-1)
    k = torch.exp(-dists / sigma / 2)
    der = (d * k[:, :, None]).sum(axis=0) / sigma
    return k.mm(scores_x) - der