Update references to JAX's GitHub repo

JAX has moved from https://github.com/google/jax to https://github.com/jax-ml/jax

PiperOrigin-RevId: 702886844

Author: jakeharmon8

README.md

@@ -98,7 +98,7 @@ pip install -r vit_jax/requirements.txt
 pip install -r vit_jax/requirements-tpu.txt
 ```
 
-For newer versions of [JAX](https://github.com/google/jax), follow the instructions
+For newer versions of [JAX](https://github.com/jax-ml/jax), follow the instructions
 provided in the corresponding repository linked here. Note that installation
 instructions for CPU, GPU and TPU differs slightly.