How to improve speed and efficiency in ODE workloads with JAX

[nicholasjng]

Hello Jax community! I started using the library a few weeks ago and I'm really enjoying it so far. :--)

I want to contribute to the project in the near future, which is why I am practicing and trying to get some hands-on experience with it now by porting a personal numpy project of mine about ODEs to Jax.

But after porting the first two solvers today (Euler + Heun methods, so very vanilla ones) and running some benchmarks, I noticed a big drop in performance compared to numpy after adopting Jax.

Right now, my port is very unsophisticated - so far, I really only switched import numpy as np to import jax.numpy as jnp and replaced direct array assignments with the array.at().set() paradigm. And yet, I'm seeing significant reductions in iterations per second for a simple exponential decay model (from 42k -> 160it/s with Heun's method).

Some things which I suspect are killing performance right now:

• I am still using raw Python loops instead of Jax's fori_loop and while_loop.
• I have not yet used jit to compile anything (and I am not yet sure about where to, either).
• I think the constant use of .at().set() is also bad (since from what I understood, that allocates a new array because Jax arrays are immutable).

So obviously, I am still doing a lot of things wrong. Therefore, I was hoping to have a discussion about some best practices for Jax (and maybe ODE solving with Jax, specifically) that I can use to write more efficient Jax code in the future.

All the best,
Nicholas

PS: I am aware of the odeint example in the jax.experimental codebase, and I have studied the code quite a bit. But I wanted to extend that example to be more dynamic, by enabling things such as solver selection, step size control algorithm selection, logging some metrics on a per-step basis etc.

[C-J-Cundy]

Awesome to hear you're enjoying exploring jax!
I should preface this by saying that I'm not an expert in ODEs, but as far as general tips with using jax goes:

Jax is really designed to be used under jit as far as possible. The un-jitted jax equivalents of numpy functions are generally slower than their numpy counterparts, since numpy is optimized that way. Plus, you need to jit to use any accelerators like a GPU, TPU etc.
As a rule of thumb, try to jit the top-most function of your algorithm. If you jit at the top level, the compiler can see everything underneath and make nice optimizations. When I first started using jit I had a lot of issues, but using static_argnums can help make this more manageable, although keep in mind that your function will be recompiled whenver the static argnums arguments are different.

Secondly, you'll want to avoid python loops and control flow with jax. Raw python loops are slow if they're not under a jit, and if they are in a jitted function then this can blow up the compilation time since python loops are statically unrolled by the compiler, which generally results in superlinear compilation time with the length of the loop. Replacing for loops with scan, fori_loop etc will really cut down on compilation time!

Using .at().set() is not very fast when not under a jit, and you probably want to use things like jax scans and vmaps if possible (though I guess vmaps are less applicable for ODEs which are quite sequential in nature rather than parallelisable).

Could be worth checking this repository out, which has quite a few ODE solvers implemented in JAX

[nicholasjng]

Thank you for the reply! My journey so far is definitely interesting, as I was able to achieve substantial gains even with unjitted jax.numpy over my old numpy code (which must have been really terrible to begin with, given what you wrote about speed comparisons between the two in your first paragraph).

So far, porting has been a relatively seamless experience, save from the technical debt I incurred by some of my own design choices, which had little to do with numpy. But now I'm really excited about running this code on my GPU, at least eventually when I can get jit to work with it.

All in all, I think I got some good directions for improvement already - the rest will come with more development experience, I guess. So I'll mark it as resolved. Again, thanks for your input!