Forward mode "adjoint"

[johannahaffner]

Here you go! This is the pragmatic solution, without support or test coverage for integer inputs and only a small comment explicating that forward mode is not really an adjoint, even though its diffrax interface is that of an adjoint.

Changes with respect to the last PR:

• renamed to ForwardMode everywhere
• add a sentence to the docstring that explains that this is not really an adjoint, but keep inheriting from AbstractAdjoint
• remove stub for "forward gradient with int" from test_adjoint.py and explain that since JAX does not offer this option, we're not writing our own workaround to test it either

On the last point: if I understood this correctly, then supporting this would entail writing a gradient-computation directly from a JVP with custom "unit pytrees". This is somewhat annoying for mixed array and non-array types.
I'm happy to try again if computing gradients with respect to integer elements of a PyTree is an expected use case (maybe arising from composed/layered transformations of a solve) that requires test coverage.

Earlier comments here.

(This is now rebased on main.)

[patrick-kidger]

Nita aside this LGTM! Once you fix things up I'll merge this into a new dev branch.

[patrick-kidger]

I might specify what an adjoint is here, to make this clear.

[johannahaffner]

Gave it a shot!

[johannahaffner]

Completed the above fixes :)

I also added references to ForwardMode to the docstrings of the other adjoints, where they point to DirectAdjoint for forward-mode automatic differentiation. I'm recommending it for the case where only forward-mode automatic differentiation is required, but I'm not sure if I should hedge/tone down my statement that ForwardMode is more efficient than DirectAdjoint. That is - I'm not sure if this statement holds in general. What do you think?

While going though the errata raised by the other adjoints I was also wondering if there is a case ForwardMode cannot handle, such as the UnsafeBrownianPath for which some of the other adjoints error out.

[johannahaffner]

I benchmarked these a little bit by computing the Jacobian of the output ys for a variable number of states, once with one parameter for each state and once with a single parameter for all states. (So the Jacobians had shape (num_timepoints, num_states, num_states) in the first case and (num_timepoints, num_states, 1) in the second case.)

I did not include any couplings between states, but this roughly makes what I have seen on my own data (with coupled ODEs), where the DirectAdjoint takes about 1.5x as long as ForwardMode.

Equal number of states and parameters, uncoupled ODE

Largest Jacobian (1024, 512, 512). DirectAdjoint takes ~2.3x as long at high and low end.

One parameter for all states, uncoupled ODE

Largest Jacobian (4096, 512, 1). DirectAdjoint takes ~1.15x as long at high end, ~2.4x as long at low end

[patrick-kidger]

LGTM! And merged :)

As for DirectAdjoint -- yup, it's totally expected that this will be slower than ForwardMode. Honestly DirectAdjoint is kind of a terrible choice, it only exists because of the very occasional need to support both forward and reverse mode through the same solve. (Hopefully one day JAX implements jvp-of-custom_vjp.)