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

agent_based_epidemic_sim/learning/covid_risk_score_sim_learning.ipynb

@@ -3116,7 +3116,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/google/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
+            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/jax-ml/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
           ]
         },
         {
@@ -3986,7 +3986,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/google/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
+            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/jax-ml/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
           ]
         },
         {
@@ -4169,7 +4169,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/google/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
+            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/jax-ml/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
           ]
         },
         {
@@ -10498,7 +10498,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/google/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
+            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:77: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[array(seq)]` instead of `arr[seq]`. In the future this will result in a TypeError. See https://github.com/jax-ml/jax/issues/4564 for discussion of why this type of indexing is being deprecated.\n"
           ]
         },
         {