Some experiments comparing CPU vs Numba vs JAX: https://github.com/google/jax Dependencies can be installed via pip install -r requirements.txt