Comment by PartiallyTyped

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/...