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

oryx/examples/notebooks/a_tour_of_oryx.ipynb

@@ -93,7 +93,7 @@
         "id": "F_n9c7K3xdKQ"
       },
       "source": [
-        "Oryx is an experimental library that extends [JAX](https://github.com/google/jax) to applications ranging from building and training complex neural networks to approximate Bayesian inference in deep generative models. Like JAX provides `jit`, `vmap`, and `grad`, Oryx provides a set of **composable function transformations** that enable writing simple code and transforming it to build complexity while staying completely interoperable with JAX.\n",
+        "Oryx is an experimental library that extends [JAX](https://github.com/jax-ml/jax) to applications ranging from building and training complex neural networks to approximate Bayesian inference in deep generative models. Like JAX provides `jit`, `vmap`, and `grad`, Oryx provides a set of **composable function transformations** that enable writing simple code and transforming it to build complexity while staying completely interoperable with JAX.\n",
         "\n",
         "JAX can only safely transform pure, functional code (i.e. code without side-effects). While pure code can be easier to write and reason about, \"impure\" code can often be more concise and more easily expressive.\n",
         "\n",