Skip to content

Instantly share code, notes, and snippets.

@apaszke
Last active April 3, 2024 03:40
Show Gist options
  • Save apaszke/226abdf867c4e9d6698bd198f3b45fb7 to your computer and use it in GitHub Desktop.
Save apaszke/226abdf867c4e9d6698bd198f3b45fb7 to your computer and use it in GitHub Desktop.
import torch
def jacobian(y, x, create_graph=False):
jac = []
flat_y = y.reshape(-1)
grad_y = torch.zeros_like(flat_y)
for i in range(len(flat_y)):
grad_y[i] = 1.
grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
jac.append(grad_x.reshape(x.shape))
grad_y[i] = 0.
return torch.stack(jac).reshape(y.shape + x.shape)
def hessian(y, x):
return jacobian(jacobian(y, x, create_graph=True), x)
def f(x):
return x * x * torch.arange(4, dtype=torch.float)
x = torch.ones(4, requires_grad=True)
print(jacobian(f(x), x))
print(hessian(f(x), x))
@maryamaliakbari
Copy link

I think this has now been added to recent versions of torch's autograd module. Maybe look at the examples here

Right. I checked it. When I use this method I am getting multiple errors. I am looking for an example or similar code to see how the implementation is done.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment