How to set a function's vjp to be identical to another's?

[ahuang314]

Suppose I have two different implementations of the same function, f1 and f2. For straightforward function evaluations, it is always better to use f1, but for calculating the derivative, it is always better to use f2. Is there a quick and dirty way of setting it so that when autodifferentiating through f1, it actually just autodifferentiates through f2 instead?

Here is what I've tried:

@custom_vjp
@jit
def f1(t, z):
    return something
    
@jit
def f1_fwd(t, z):
    y = f1(t, z)
    jac_t = jax.jacobian(f2, argnums=0)(t, z)
    jac_z = jax.jacobian(f2, argnums=1)(t, z)
    return y, (jac_t, jac_z)

@jit
def f1_bwd(res, g):
    jac_t, jac_z = res
    return (g @ jac_t, g @ jac_z)

f1.defvjp(f1_fwd, f1_bwd)

However this still seems to be slower than autodifferentiating through f2 directly. Additionally, I seem to run into some assertion error when z is complex, even after setting holomorphic=True in the jacobian call.

[dfm]

I think that you might be better off using something like:

@jax.custom_jvp
def f1(*args):
  return ...

@f.defjvp
def f1_jvp(primals, tangents):
  return jax.jvp(f2, primals, tangents)

or if you really don't ever want to use f2 for the primals:

@f.defjvp
def f1_jvp(primals, tangents):
  _, out_tangents = jax.jvp(f2, primals, tangents)
  return f1(*primals), out_tangents

The first one will have exactly the same computational performance as f2 when differentiated, whereas the second one might perform some extra work, since the primals aren't used from f2's forward pass.

If this doesn't work, please feel free to post a minimal reproducer with the errors your seeing!

[ahuang314]

It works perfectly. Thank you!