Hi,
In my use case, I need to compute the diagonal of the Hessian for 3D tensors.
The code below illustrates what I need to compute. It works, however, it does not scale well if one increases the dimension of the tensor.
Are there faster way to achieve the same computation?
I am on torch=1.10.0.
from time import time

import torch
from torch._vmap_internals import _vmap as vmap

def grad(outputs, inputs, **kwargs):
    """
    Compute the trace of the jacobian of a 3D tensor.

    grad_{ijk} = d (outputs_{ijk}) / d (inputs_{ijk})

    See: https://discuss.pytorch.org/t/jacobian-functional-api-batch-respecting-jacobian/84571/7?u=amerlo94
    """

    shape = outputs.shape
    bs = shape[0]
    n = shape[1] * shape[2]

    outputs = outputs.view(bs, n)
    outputs = outputs.sum(axis=0)
    grad_outputs = torch.eye(n, dtype=outputs.dtype, device=outputs.device)

    def get_vjp(v):
        return torch.autograd.grad(outputs, inputs, v, **kwargs)[0]

    vjp = vmap(get_vjp)
    grad = vjp(grad_outputs)

    return grad.T.view(shape)


n = 16
b = torch.randn(1, n, n)


def f(x):
    return x ** 2 * b


bs = 128
iterations = 1000

x = torch.rand(bs).requires_grad_(True)
x = x.view(-1, 1, 1)
y = f(x)

#  Check gradients
jac = grad(y, x, create_graph=True)
hes = grad(jac, x)
assert torch.allclose(jac, 2 * x * b)
assert torch.allclose(hes, 2 * b)

#  Time gradients
t0 = time()
for _ in range(iterations):
    jac = grad(y, x, create_graph=True)
    grad(jac, x)

tottime = time() - t0

time_per_batch = tottime / iterations / bs
print(f"Time per batch: {time_per_batch * 1e6:.2f} us")