import torch


def sum_hessian(loss, hidden):
    """Given scalar tensor 'loss' and [b, d] batch of hidden activations 'hidden', computes [d, d] size sum of hessians
    for all entries in the batch.

    We make use of the fact that grad^2(f) = grad(grad(f)) and that sum of grads = grad of sum. So, sum(grad^2(f)) is
    computed as grad(sum(grad(f))).
    """
    d = hidden.size(1)
    grad = torch.autograd.grad(loss, hidden, retain_graph=True, create_graph=True)[0]
    sum_grad = torch.sum(grad, dim=0)
    hessian = hidden.new_zeros(d, d)
    for i in range(d):
        hessian[i, :] = torch.autograd.grad(sum_grad[i], hidden, retain_graph=True)[0].sum(dim=0)
    return hessian