import torch


def ssm_vr_loss(energy_model, x, n_slices=1):
    """SSM-VR (variance reduction) loss from
    Sliced Score Matching: A Scalable Approach to Density and Score Estimation

    The loss is computed as
    s = -dE(x)/dx
    loss = vT*(ds/dx)*v + 1/2*||s||^2

    Args:
        x (torch.Tensor): input samples
        n_slices (torch.Tensor): number of slices, default=1

    Returns:
        SSM-VR loss
    """
    x = x.unsqueeze(0).expand(n_slices, *x.shape)  # (n_slices, b, ...)
    x = x.contiguous().view(-1, *x.shape[2:])  # (n_slices*b, ...)
    x = x.requires_grad_()
    score = energy_model.score(x)  # (n_slices*b, ...)
    v = torch.randn((n_slices,) + x.shape, dtype=x.dtype, device=x.device)
    v = v.view(-1, *v.shape[2:])  # (n_slices*b, ...)
    sv = torch.sum(score * v)  # ()
    loss1 = torch.norm(score, dim=-1) ** 2 * 0.5  # (n_slices*b,)
    gsv = torch.autograd.grad(sv, x, create_graph=True)[0]  # (n_slices*b, ...)
    loss2 = torch.sum(v * gsv, dim=-1)  # (n_slices*b,)
    loss = (loss1 + loss2).mean()  # ()
    return loss