Jax conversion of custom jvps and vjps (and choice of forward or backward mode diff).

[LawrenceMMStewart]

Hello,

I have a couple of questions relating to custom jvps and vjps, as well as forward and reverse mode differentiation in Jax.

Suppose I have a chain of function compositions, and wish to invoke Jax's auto-diff. An example of this could be calculating the loss of a neural networks prediction for a batch of data. For illustrative purposes, lets assume that the neural network consists of conventional layers (i.e. those available in Jax libraries e.g. nn.Linear from Flax) but also contains some custom user-defined layers, some of which have custom jvps and some of which have custom vjps.

• Firstly, as a sanity check, am i correct in assuming that when calling jax.grad or jax.value_and_grad, Jax will have to choose to use only one of forward / reverse mode differentiation (i.e. it is not possible for some layers to use jvps in their computation and others to use vjps)?

• If the above is indeed true, then Jax must convert any custom defined jvps into vjps if using reverse-mode differentiation or vice versa if using forward mode differentiation. How does Jax decide whether to do forward or reverse mode diff (e.g. does it just calculate the total number of Ops for both forward and reverse mode using the dimensions of all traced arrays and pick the minimum) ?

• If one is defining a custom layer for a model, it might seem easier in terms of implementations to use jax.custom_jvp instead of jax.custom_vjp. This is because in the implementation all the intermediary calculations for the forward are available for calculating the jvp, so one can be sure that calculations are not being duplicated on the backward. Furthermore, for reverse-mode auto-diff, it seems that Jax is decomposing jvps (using linearization and transposition) to perform vjps, so implementing everything as a custom_jvp seems to fit in with this philosophy. With both of the above reasons in mind, is it always preferential to use custom_jvp even when we know reverse mode diff is going to be used (i.e. when would one prefer to use custom_vjp instead)?

• Can a user force either forward or reverse mode differentiation (or is this always decided by Jax)?

Many thanks for your help, please do not hesitate to let me know if any of the questions are unclear.

[dionhaefner]

jax.grad and jax.value_and_grad always use reverse-mode AD. Forward-mode AD is exposed as jax.jacfwd and jax.jvp.

(See e.g. here: grad doesn't work if VJP / reverse-mode isn't defined on a primitive.)

[LawrenceMMStewart]

Great, many thanks for clarifying!