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

Author: jakeharmon8

README.md

@@ -19,7 +19,7 @@ implementation of the [K-FAC] optimizer and curvature estimator.
 
 KFAC-JAX is written in pure Python, but depends on C++ code via JAX.
 
-First, follow [these instructions](https://github.com/google/jax#installation)
+First, follow [these instructions](https://github.com/jax-ml/jax#installation)
 to install JAX with the relevant accelerator support.
 
 Then, install KFAC-JAX using pip:
@@ -219,6 +219,6 @@ and the year corresponds to the project's open-source release.
 
 
 [K-FAC]: https://arxiv.org/abs/1503.05671
-[JAX]: https://github.com/google/jax
+[JAX]: https://github.com/jax-ml/jax
 [Haiku]: https://github.com/google-deepmind/dm-haiku
 [documentation]: https://kfac-jax.readthedocs.io/

docs/index.rst

@@ -3,7 +3,7 @@
 KFAC-JAX Documentation
 ======================
 
-KFAC-JAX is a library built on top of `JAX <https://github.com/google/jax>`_ for
+KFAC-JAX is a library built on top of `JAX <https://github.com/jax-ml/jax>`_ for
 second-order optimization of neural networks and for computing scalable
 curvature approximations.
 The main goal of the library is to provide researchers with an easy-to-use
@@ -16,7 +16,7 @@ Installation
 
 KFAC-JAX is written in pure Python, but depends on C++ code via JAX.
 
-First, follow `these instructions <https://github.com/google/jax#installation>`_
+First, follow `these instructions <https://github.com/jax-ml/jax#installation>`_
 to install JAX with the relevant accelerator support.
 
 Then, install KFAC-JAX using pip::

examples/losses.py

@@ -104,7 +104,7 @@ def softmax_cross_entropy(
   max_logits = jnp.max(logits, keepdims=True, axis=-1)
 
   # It's unclear whether this stop_gradient is a good idea.
-  # See https://github.com/google/jax/issues/13529
+  # See https://github.com/jax-ml/jax/issues/13529
   max_logits = lax.stop_gradient(max_logits)
 
   logits = logits - max_logits