


from torch.func import jacfwd, vmap
x = torch.randn(64, 5)
jacobian = vmap(jacfwd(torch.sin))(x)
assert jacobian.shape == (64, 5, 5)



If you would like to compute the output of the function as well as
   the jacobian of the function, use the "has_aux" flag to return the
   output as an auxiliary object:



from torch.func import jacfwd
x = torch.randn(5)
def f(x):
  return x.sin()
def g(x):
  result = f(x)
  return result, result
jacobian_f, f_x = jacfwd(g, has_aux=True)(x)
assert torch.allclose(f_x, f(x))



Additionally, "jacrev()" can be composed with itself or "jacrev()"
   to produce Hessians



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



By default, "jacfwd()" computes the Jacobian with respect to the