Gather Jacobians with regards to the input in the forward-propagation step

[LeanderK]

I am currently using pytorch but wonder whether I should to switch to JAX. I want to do the following:

Let's say I have a network composed of layers: input -> layer 1 -> layer 2 -> ... -> layer n -> output

I need the Jacobian with regards to the input for each layer in a differentiable way, so that I can use it in the loss, without being too inefficient. So what I want to optimise is:

loss = criterion(network(input), targets) + alpha*sum_{layer i in layers} f(Jacobian_i(input))

I think the most straightforward and reasonably efficient way is to use forward-mode differentiation to record the Jacobians, but I wonder how this is possible? I don't see a way to record the interim jacobians during forward-propagation.