Hello
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?
Thanks!