Benchmarking Numba vs JAX

[tomwhite]

We have used Numba to accelerate CPU computations in sgkit for a long time - and to great effect. We have previously discussed trying JAX as an alternative to see how it compares.

I created a notebook to compare basic allele counting (count_call_alleles) using Numba with a JAX equivalent here: https://github.com/tomwhite/sgkit/blob/134f15c18558458f6b04607cca9915e5ca90acf6/benchmark-numba-vs-jax.ipynb

The JAX code I wrote is significantly slower than Numba in this case. It could be because I've used JAX's bincount function, rather than writing code that operates directly on the arrays - like Numba does - but I'm not sure how to do that with JAX.

[jeromekelleher]

Thanks @tomwhite that's very interesting. I think the limitations of having to use a numpy-like API rather than simple loops, and this significant performance drop on key functionality is a pretty strong argument for sticking with Numba. Numba has it's issues, but overall I'm still very positive about it.

[tomwhite]

I did have a look at using JAX's lower-level fori_loop, but I couldn't see how it could loop over array dimensions (or if that is even possible).

Also, the approach I've used for count_call_alleles won't work forhardy_weinberg_test (as @eric-czech pointed out to me), since it is sufficiently complicated to not be expressible in NumPy-like operations, and would therefore need loops.

[tomwhite]

I did have a look at using JAX's lower-level fori_loop, but I couldn't see how it could loop over array dimensions (or if that is even possible).

I managed to express count_alleles using the lower-level jax.lax.scan primitive - but it's still an order of magnitude slower than the Numba equivalent.

Notebook at https://github.com/sgkit-dev/sgkit/blob/df5f7a146ba64a0654d00bc5310e7b64c329b19e/benchmark-numba-vs-jax.ipynb

[jeromekelleher]

It's pretty obscure code too. Numba FTW!