"argnums".
Note:
 You may see this API error out with "forward-mode AD not
 implemented for operator X". If so, please file a bug report and
 we will prioritize it. An alternative is to use
 "jacrev(jacrev(func))", which has better operator coverage.

A basic usage with a R^N -> R^1 function gives a N x N Hessian:



from torch.func import hessian
def f(x):
  return x.sin().sum()
x = torch.randn(5)
hess = hessian(f)(x)  # equivalent to jacfwd(jacrev(f))(x)
assert torch.allclose(hess, torch.diag(-x.sin()))


