Hessians



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



By default, "jacrev()" computes the Jacobian with respect to the
   first input. However, it can compute the Jacboian with respect to a
   different argument by using "argnums":



from torch.func import jacrev
def f(x, y):
  return x + y ** 2
x, y = torch.randn(5), torch.randn(5)
jacobian = jacrev(f, argnums=1)(x, y)
expected = torch.diag(2 * y)
assert torch.allclose(jacobian, expected)



Additionally, passing a tuple to "argnums" will compute the
   Jacobian with respect to multiple arguments



from torch.func import jacrev
def f(x, y):
  return x + y ** 2
x, y = torch.randn(5), torch.randn(5)
jacobian = jacrev(f, argnums=(0, 1))(x, y)
expectedX = torch.diag(torch.ones_like(x))


