Earlier quoted context omitted.
I've been using JAX ( https://jax.readthedocs.io/en/latest/ ) for scientific computing in general (in particular MCMC algorithms), as it's really fast. Even on a CPU you get massive speedups compared to numpy (can be up to 2 or 3 orders of magnitude faster in some cases). The main selling point of the library is automatic differential and compilation to XLA, but I've been using it even when I don't need gradients, as…
Have you tried numba+numpy? In my experience, it is much faster than Jax and can compile to cuda. It's not caveat free, but it also removes the hustle of labeling arrays as donated in Jax. You may find this interesting https://github.com/scikit-hep/iminuit/blob/develop/tutorial/...
With Jax you can write native for loops that can also be jitted (I imagine you can also do this in numba?); this can then be really fast. Though in that case you would have to write the optimisation algorithm yourself which is not always practical!
Another big speedup in Jax is due to vmap/pmap, which allow to vectorise/parallelise computation. For example you can build a massive gram matrix really quickly using vmap.
Another point: Jax can also run on GPU (like numba :) ) without having to rewrite anything.