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

clu/profiler.py

@@ -31,8 +31,9 @@ def start(logdir: str, options=None):
   if options is not None:
     raise NotImplementedError(
         "'options' not supported by clu.profiler.start(). Please file an issue "
-        "at https://github.com/google/jax/issues requesting profiler option "
-        "support if you need this feature.")
+        "at https://github.com/jax-ml/jax/issues requesting profiler option "
+        "support if you need this feature."
+    )
   if logdir is None:
     raise ValueError("Must specify logdir where profile should be written!")
   jax.profiler.start_trace(logdir)

clu_synopsis.ipynb

@@ -16,7 +16,7 @@
         "https://github.com/google/CommonLoopUtils\n",
         "\n",
         "This package is usually used with\n",
-        "[JAX](https://github.com/google/jax)\n",
+        "[JAX](https://github.com/jax-ml/jax)\n",
         "/\n",
         "[Flax](https://github.com/google/flax)\n",
         "ML projects, but it can also be used with other ML frameworks.\n",