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

Author: jakeharmon8

clu/profiler.py

@@ -31,7 +31,7 @@ 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 "
+        "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!")

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",