Say that I want to calculate the output of a function and its Jacobian. One way to do it would be to do the following
import torch
x = torch.rand(5, requires_grad=True)
net = torch.nn.Linear(5, 2)
y = net(x)

jac = torch.autograd.functional.jacobian(net, x)

however, this means that I’ll have to do two forward passes. It feels like there should be a way to do the following
jac = torch.autograd.functional.jacobian(lambda x: y, x)

but then jac is all zeros. Is there a way to do this?
The end goal is to use the Jacobian with respect to the network’s activations in a regularization term, but this is the step I’m bothered about.