Hutchinson's trick and batch grad_outputs for autograd.grad

I am trying to use Hutchinson’'s trick to compute an estimator of the trace of a hessian matrix.
To make this estimator precise, I need to run the same backward pass several times with different stochastic vectors v that are used for vector-jacobian product.

The code using a for loop would look like this:

def hutch_loop(loss, params, niter, v):
    g = autograd.grad(loss, params, create_graph=True)
    out = 0.0
    for i in range(niter):
        vsplit = [torch.randn_like(p) for p in params]
        Hv = autograd.grad(g, params, vsplit, retain_graph=True)
        out += group_product(Hv, vsplit)
    return out/niter

def group_product(xs, ys):
    return sum([torch.sum(x * y) for (x, y) in zip(xs, ys)])

Ideally, I would like to be able to do this operation in batch mode, alleviating the need to compute as many backward passes than the number of samples i need to draw. Is there any way to do this efficiently?



I’m afraid that right now this is the best that you can do.
But we are actively working on changing that. If all goes according to plan, we should have a batched version for the fall of this year.