Some experiments comparing CPU vs Numba vs JAX: https://github.com/google/jax

Dependencies can be installed via

pip install -r requirements.txt