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

Author: jakeharmon8

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]
 
 

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]
 
 

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