From 129ff3aa9959df27e13b90b608e3f471d55af7cb Mon Sep 17 00:00:00 2001 From: Jake Harmon Date: Wed, 4 Dec 2024 15:46:04 -0800 Subject: [PATCH] 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 --- scenic/projects/baselines/bert/train_utils.py | 2 +- scenic/projects/layout_denoise/train_utils.py | 2 +- scenic/projects/token_learner/model.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/scenic/projects/baselines/bert/train_utils.py b/scenic/projects/baselines/bert/train_utils.py index 5dc0d6471..616ee0716 100644 --- 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 index 3ee2d801b..21d3cbdf3 100644 --- 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 index 072dea9cb..03ce1af18 100644 --- 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