JAX slow compared to numba for certain calculations

Hi there,

I recently coded up a version of the Glicko rating system and found that for the particular calculations involved, using jit decorators in JAX runs slowly, while the same strategy works well in numba. I put together a benchmark and notes here:

https://github.com/martiningram/jax_vs_numba_glicko

Naively using jit decorators runs the benchmark in 20s in JAX, compared to 121ms in numba, so clearly that’s not a good strategy with JAX. @mattjj asked me to share this here, since while this is not a particular workload you have been looking at with JAX, it could be good to understand why it is slow and perhaps how to improve things. I hope this is useful!

[jekbradbury]

Somewhat off topic for the JAX issue tracker, but this is a pretty ideal case for the Julia compiler, which can run calculate_approximate_likelihood in about 400ns, or a little less time than it takes to generate_random_observations: https://github.com/martiningram/jax_vs_numba_glicko/pull/1. Since all three of XLA/Numba/Julia are using LLVM, all this really means is that JAX/XLA has even more room to grow on heavily-scalar applications like this one; I don’t see any reason why we wouldn’t (eventually, in principle) be able to generate ~the same code Julia does here.

[jakevdp]

Closing, as I don’t think there is any action item here. Ongoing work on dynamic shapes may improve the situation with respect to shape-agnostic JIT compilation, but that’s tracked elsewhere.