Upgrade JAX dependency

I'm curious how much work would be involved to upgrade the JAX dependency from [JAX 0.6.2](https://github.com/keras-team/keras/blob/master/requirements-jax-cuda.txt#L11) to JAX 0.8. I'm trying to build from source on macOS, but JAX 0.6 requires Bazel 6, which doesn't build with Apple Clang 17.