It wasn't me!

Limits the maximum supported jax version to 0.4.6 due to the new jax.jit implementation triggering a bug I have not been able yet to fix.