Reconfigure JAX in Colab to use latest image.

PiperOrigin-RevId: 721679162
Change-Id: Ib15e5ff170687b2ef2fe2af206ebce4d9803b928

Author: andrewlkd

gencast_mini_demo.ipynb

@@ -75,6 +75,21 @@
         "%pip install --upgrade https://github.com/deepmind/graphcast/archive/master.zip"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HH8S9fND3KRB"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Reconfigure jax if running on TPU.\n",
+        "\n",
+        "# This is required due to outdated jax and libtpu versions in Colab TPU images.\n",
+        "%pip uninstall -y libtpu libtpu-nightly\n",
+        "%pip install -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,