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
google · Dec 11, 2024 · a82cfde · a82cfde
1 parent 8daea3c
commit a82cfde
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -63,7 +63,7 @@ jax.grad(loss_fn)(scores, labels)  # computes gradients w.r.t. scores.
 
 ## Installation
 
-See https://github.com/google/jax#installation for instructions on installing JAX.
+See https://github.com/jax-ml/jax#installation for instructions on installing JAX.
 
 We suggest installing the latest stable version of Rax by running:
 

diff --git a/docs/index.rst b/docs/index.rst
@@ -26,7 +26,7 @@ Rax is a Learning-to-Rank (LTR) library built on top of JAX.
 Installation
 ------------
 
-See https://github.com/google/jax#pip-installation for instructions on
+See https://github.com/jax-ml/jax#pip-installation for instructions on
 installing JAX.
 
 We suggest installing the latest stable version of Rax by running::