jax

3a8f31aa - Update the JAX version to 0.5.0.

Update the JAX version to 0.5.0. This is because of the breaking change to PRNG key semantics, and the version follows JAX's new effver versioning scheme (https://jax.readthedocs.io/en/latest/jep/25516-effver.html).