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: 702886845
google-research · Dec 5, 2024 · 129ff3a · 129ff3a
1 parent 96db8b0
commit 129ff3a
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/scenic/projects/baselines/bert/train_utils.py b/scenic/projects/baselines/bert/train_utils.py
@@ -34,7 +34,7 @@
 import sklearn.metrics
 
 # JAX team is working on type annotation for pytree:
-# https://github.com/google/jax/issues/1555
+# https://github.com/jax-ml/jax/issues/1555
 PyTree = Union[Mapping[str, Mapping], Any]
 
 

diff --git a/scenic/projects/layout_denoise/train_utils.py b/scenic/projects/layout_denoise/train_utils.py
@@ -32,7 +32,7 @@
 from scenic.train_lib_deprecated import train_utils
 
 # JAX team is working on type annotation for pytree:
-# https://github.com/google/jax/issues/1555
+# https://github.com/jax-ml/jax/issues/1555
 PyTree = Union[Mapping[str, Mapping], Any]
 
 

diff --git a/scenic/projects/token_learner/model.py b/scenic/projects/token_learner/model.py
@@ -39,7 +39,7 @@
 from scenic.projects.vivit import model_utils as vivit_model_utils
 
 # JAX team is working on type annotation for PyTree:
-# https://github.com/google/jax/issues/1555
+# https://github.com/jax-ml/jax/issues/1555
 Array = Union[jnp.ndarray, np.ndarray]
 PyTree = Any