Is there a way to calculate cotangents with respect to particular outputs?

[NeilGirdhar]

Suppose you have some opaque, expensive function:

def f(v, w) -> Scalar, Scalar:
  return loss_a, loss_b

Given some v, w, how do you find the derivative of loss A with respect to v, and the derivative of loss B with respect to w? The naive approach is to do:

def f_a(v_variable):
  return f(v_variable, w)[0]

def f_b(w_variable):
  return f(v, w_variable)[1]

v_bar = grad(f_a)(v)
w_bar = grad(f_b)(w)

(This is just typed by intuition, so apologies if I've made an error.)

Is it possible to only call f once? Perhaps Jax has some way of keeping track of multiple cotangents on variables?

[mattjj]

Good question, as always!

One way to reduce redundant work here (without relying on jax.jit to do common-subexpression elimination) is to use jax.vjp directly (around which grad is a thin wrapper):

(loss_a, loss_b), f_vjp = jax.vjp(f, v, w)
v_bar, _ = f_vjp(jnp.ones_like(loss_a), jnp.zeros_like(loss_b))
_, w_bar = f_vjp(jnp.zeros_like(loss_a), jnp.ones_like(loss_b))

That will only run the forward pass once (on the line which calls jax.vjp) and then run two separate backward passes.

Another variant would be just to call jax.jacrev(f)(v, w), then extract the entries of the result that you want (corresponding to v-input-a-output and w-input-b-output). That'll also run the forward pass just once. The memory usage is a bit different between these two approaches; the latter is basically just using a jax.vmap to batch together what are two separate calls to the f_vjp function in the former.

WDYT?

[mattjj]

I made an edit to fix a typo!

[NeilGirdhar]

Oh of course! Thanks for being helpful as always!