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: 702886818

README.md

@@ -2,7 +2,7 @@
 
 ![PyPI version](https://img.shields.io/pypi/v/jax-triton)
 
-The `jax-triton` repository contains integrations between [JAX](https://github.com/google/jax) and [Triton](https://github.com/openai/triton).
+The `jax-triton` repository contains integrations between [JAX](https://github.com/jax-ml/jax) and [Triton](https://github.com/openai/triton).
 
 Documentation can be found [here](https://jax-ml.github.io/jax-triton).
 

docs/index.md

@@ -1,6 +1,6 @@
 # JAX-Triton documentation
 
-JAX-Triton is a repository containing containing integrations between [JAX](https://github.com/google/jax)
+JAX-Triton is a repository containing containing integrations between [JAX](https://github.com/jax-ml/jax)
 and [Triton](https://github.com/openai/triton).
 
 JAX is a Python library for accelerated numerical computing and Triton is a Python library and compiler for writing custom GPU kernels.
@@ -20,7 +20,7 @@ For example you could run:
 ```bash
 $ pip install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
-Check out the [JAX installation guide](https://github.com/google/jax#pip-installation-gpu-cuda) for details.
+Check out the [JAX installation guide](https://github.com/jax-ml/jax#pip-installation-gpu-cuda) for details.
 
 ### Installation at HEAD