Why not use Numba? In my mental model they achieve the same win, generating GIL-free, bytecode-level perf.

I'm not super familiar with JAX though.