jax.jacrev performance: custom_vjp vs custom_jvp

[antalszava]

Hi everyone, assuming we want to get the Jacobian of a differentiable function f:R^N->R^M using a custom rule, how does the performance of

1. jax.jacrev with custom_vjp compare to
2. jax.jacrev with custom_jvp?

I've come across the note on Decomposing reverse-mode automatic differentiation - I'd be keen to learn more about how the decomposition (if at all) affects performance.

Does the performance difference between 1. and 2. (assuming it's significant) depend on:

• Whether N << M or N >> M;
• The type of computation performed in the custom_vjp/custom_jvp rule;
• Jitting;
• Something else?

Let me know if it would be helpful to share more about the f or it's potential VJP/JVP rule that's being considered. The case of N being big (e.g., a couple of thousand or more) and M=1 is of interest for example, but curious about general insights too.

[mattjj]

Thanks for the question!

Whether you use custom_jvp or custom_vjp need not affect performance at all. In fact you can "spell" the exact same computation either way in most cases.

I say "in most cases" as a hedge because there are cases where they might differ. Here is a non-exhaustive list which comes to mind:

• With custom_vjp you can more directly control what the "residuals" are, i.e. what values are saved from the forward pass for use on the backward pass. With custom_jvp you're leaving the residuals up to JAX's partial evaluation machinery. Those can be controlled too (e.g. with jax.remat/jax.checkpoint) but it's not "built in" in the way it is with custom_vjp. Choice of residuals can affect performance. But it's mainly important for when there isn't a jit outside your jacrev; if the whole thing is under a jit, so long as loops like jax.lax.scan aren't involved, XLA will likely perform optimizations which amount to choosing the residuals for you.
• Related to the first bullet, for some computations you might need to control the choice of residuals, for example if f is an ODE integrator implemented in terms of (unbounded trip count) while_loops. In effect, for these kinds of computations you currently must use custom_vjp and not custom_jvp to get reverse-mode AD. This actually gets at a missing piece in the custom AD tools currently on offer in JAX.

I wouldn't expect the size of N or M to matter.

WDYT?

[antalszava]

This is great, thank you for the fast and detailed explanation! 🎉