From e84959cac07eb4d8e6aca12c38cfa8e297b39d18 Mon Sep 17 00:00:00 2001
From: Jake Harmon <jakeharmon@google.com>
Date: Thu, 5 Dec 2024 04:04:35 -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: 703059900
---
 README.md                  | 6 +++---
 docs/index.rst             | 2 +-
 fedjax/core/dataclasses.py | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 5a8c893..597b1cc 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,7 @@ pip install --upgrade jax jaxlib  # CPU-only version
 ```
 
 For other devices (e.g. GPU), follow
-[these instructions](https://github.com/google/jax#installation).
+[these instructions](https://github.com/jax-ml/jax#installation).
 
 Then, install FedJAX from PyPI:
 
@@ -112,7 +112,7 @@ To cite this repository:
 *   https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html
 *   https://dm-haiku.readthedocs.io/en/latest/
 
-[JAX]: https://github.com/google/jax
+[JAX]: https://github.com/jax-ml/jax
 [Haiku]: https://github.com/deepmind/dm-haiku
-[Stax]: https://github.com/google/jax/blob/main/jax/example_libraries/stax.py
+[Stax]: https://github.com/jax-ml/jax/blob/main/jax/example_libraries/stax.py
 [Optax]: https://github.com/deepmind/optax
diff --git a/docs/index.rst b/docs/index.rst
index e2bcf2b..e53ae1c 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -2,7 +2,7 @@
 FedJAX documentation
 ====================
 
-FedJAX is a library for developing custom `Federated Learning (FL) <https://ai.googleblog.com/2017/04/federated-learning-collaborative.html>`_ algorithms in `JAX <https://github.com/google/jax>`_. FedJAX prioritizes ease-of-use and is intended to be useful for anyone with knowledge of NumPy.
+FedJAX is a library for developing custom `Federated Learning (FL) <https://ai.googleblog.com/2017/04/federated-learning-collaborative.html>`_ algorithms in `JAX <https://github.com/jax-ml/jax>`_. FedJAX prioritizes ease-of-use and is intended to be useful for anyone with knowledge of NumPy.
 
 .. toctree::
 	:maxdepth: 1
diff --git a/fedjax/core/dataclasses.py b/fedjax/core/dataclasses.py
index 0dec97f..3b9bc64 100644
--- a/fedjax/core/dataclasses.py
+++ b/fedjax/core/dataclasses.py
@@ -19,7 +19,7 @@
 
 # Forked and slimmed down from
 # https://flax.readthedocs.io/en/latest/_modules/flax/struct.html#dataclass
-# https://github.com/google/jax/issues/2371
+# https://github.com/jax-ml/jax/issues/2371
 def dataclass(clz: type):
   """Creates a dataclass which can be passed to functional transformations."""
   data_clz = dataclasses.dataclass(frozen=True)(clz)