Understand diffrax.diffeqsolve

[zhengqigao]

Hi,

Thanks for the nice package. I have a question regarding the implementation of diffrax.diffeqsolve. Specifically, I want to know a bit more details about the ODE solver, especially on what are the key factors driving diffeqsolve to be faster. In a nutshell, can you list several bulletpoints on the code implementation optimizations that have been done to accelerate ODE solving in diffrax?

[patrick-kidger]

There's been a lot of JAX-specific tricks that have gone into this: for example knowing when the compiler will want to make a copy of a buffer, and avoiding those cases. Or knowing that vmap-of-cond becomes jnp.where, and likewise knowing to avoid that.

I think one interesting trick here is the way we make sure to be vmap-friendly. That is, if we do vmap(diffeqsolve), we'd like the result to be fast. The key trick here is that we don't do a loop-over-output-times (integrating in between), and instead do a loop-over-steps (outputting times as we go along). This avoids having a double-while loop, with the inner loop sitting and waiting at every output time until every batch element is done.

Another trick that comes to mind is the use of Stumm-Walther-Wang-Moin checkpointing, to make autodiff fast: https://github.com/patrick-kidger/equinox/blob/main/equinox/internal/_loop/checkpointed.py

There's probably a bunch of other things too, but this is what comes to mind!