JAX Jacobian with upstream gradient?

[karan-dalal]

Hi all! Had a question regarding efficient computation of a Jacobian with an upstream gradient. I'm writing a custom backward pass using VJPs and need to compute the following.

To begin, my custom VJP receives the upstream gradient, let's label this dL/dz. I have the following function inside my custom VJP:

def f(params, token):
   z = InnerNet.apply({'params': params}, token, train=False) # (Just computes linear product)
   return z

I want to compute dL/dW, where W is passed in as the params argument. I understand one way of doing this is computing the full jacobian of dz/dW and then multiply by the upstream gradient, but was wondering if there is a more efficient process without explicitly computing the complete jacobian.

At a high-level, does the autodiff cookbook have some sort of function that models the following behavior:

Thanks!

[karan-dalal]

@jakevdp any help would be great!

[patrick-kidger]

I think _, vjp_fn = jax.vjp(f, params, ...) will do what you want. You can then pass the dL/dz cotangents to vjp_fn and it will compute the cotangents dL/dparams.

[karan-dalal]

Many thanks Patrick!