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: 702886693

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

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

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)