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

docs/index.rst

@@ -35,7 +35,7 @@ can be used to obtain the most recent version of Optax::
    pip install git+git://github.com/google-deepmind/optax.git
 
 Note that Optax is built on top of JAX.
-See `here <https://github.com/google/jax?tab=readme-ov-file#installation>`_
+See `here <https://github.com/jax-ml/jax?tab=readme-ov-file#installation>`_
 for instructions on installing JAX.
 
 

examples/contrib/differentially_private_sgd.ipynb

@@ -12,7 +12,7 @@
         "\n",
         "A large portion of this code is forked from the differentially private SGD\n",
         "example in the [JAX repo](\n",
-        "https://github.com/google/jax/blob/main/examples/differentially_private_sgd.py).\n",
+        "https://github.com/jax-ml/jax/blob/main/examples/differentially_private_sgd.py).\n",
         "\n",
         "[Differentially Private Stochastic Gradient Descent](https://arxiv.org/abs/1607.00133) requires clipping the per-example parameter\n",
         "gradients, which is non-trivial to implement efficiently for convolutional\n",

optax/_src/alias_test.py

@@ -264,7 +264,7 @@ def test_preserve_dtype(self, opt_name, opt_kwargs, dtype):
     # x = 0.5**jnp.asarray(1, dtype=jnp.int32)
     # (appearing in e.g. optax.tree_utils.tree_bias_correction)
     # are promoted (strictly) to float32 when jitted
-    # see https://github.com/google/jax/issues/23337
+    # see https://github.com/jax-ml/jax/issues/23337
     # This may end up letting updates have a dtype different from params.
     # The solution is to fix the dtype of the result to the desired dtype
     # (just as done in optax.tree_utils.tree_bias_correction).
@@ -851,7 +851,7 @@ def test_minimize_bad_initialization(self):
     chex.assert_trees_all_close(jnp_fun(optax_sol), minimum, atol=tol, rtol=tol)
 
   def test_steep_objective(self):
-    # See jax related issue https://github.com/google/jax/issues/4594
+    # See jax related issue https://github.com/jax-ml/jax/issues/4594
     tol = 1e-5
     n = 2
     mat = jnp.eye(n) * 1e4

optax/contrib/_common_test.py

@@ -301,7 +301,7 @@ def test_preserve_dtype(
     # x = 0.5**jnp.asarray(1, dtype=jnp.int32)
     # (appearing in e.g. optax.tree_utils.tree_bias_correction)
     # are promoted (strictly) to float32 when jitted
-    # see https://github.com/google/jax/issues/23337
+    # see https://github.com/jax-ml/jax/issues/23337
     # This may end up letting updates have a dtype different from params.
     # The solution is to fix the dtype of the result to the desired dtype
     # (just as done in optax.tree_utils.tree_bias_correction).