Skip to content

Instantly share code, notes, and snippets.

@Tenglon
Forked from sbarratt/torch_jacobian.py
Created November 19, 2021 08:52
Show Gist options
  • Select an option

  • Save Tenglon/862a952fb10f697da8e22f65aed4f385 to your computer and use it in GitHub Desktop.

Select an option

Save Tenglon/862a952fb10f697da8e22f65aed4f385 to your computer and use it in GitHub Desktop.
Get the jacobian of a vector-valued function that takes batch inputs, in pytorch.
def get_jacobian(net, x, noutputs):
x = x.squeeze()
n = x.size()[0]
x = x.repeat(noutputs, 1)
x.requires_grad_(True)
y = net(x)
y.backward(torch.eye(noutputs))
return x.grad.data
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment