diagonal array as the Jacobian



from torch.func import jacrev
x = torch.randn(5)
jacobian = jacrev(torch.sin)(x)
expected = torch.diag(torch.cos(x))
assert torch.allclose(jacobian, expected)



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 jacrev
x = torch.randn(5)
def f(x):
  return x.sin()
def g(x):
  result = f(x)
  return result, result
jacobian_f, f_x = jacrev(g, has_aux=True)(x)
assert torch.allclose(f_x, f(x))



"jacrev()" can be composed with vmap to produce batched Jacobians:



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



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