Upgrade cuDNN version in Colabs to support JAX 0.4.3. Add missing "Op…

…en in Colab button". PiperOrigin-RevId: 508779352
google · Feb 11, 2023 · 46306b6 · 46306b6
1 parent adaf9e2
commit 46306b6
Show file tree

Hide file tree

Showing 11 changed files with 21 additions and 13 deletions.
diff --git a/notebooks/Disentangling_Trainability_and_Generalization.ipynb b/notebooks/Disentangling_Trainability_and_Generalization.ipynb
@@ -43,7 +43,7 @@
    "outputs": [],
    "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
    ]
   },

diff --git a/notebooks/elementwise.ipynb b/notebooks/elementwise.ipynb
@@ -66,7 +66,7 @@
       ],
       "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
       ]
     },

diff --git a/notebooks/empirical_ntk_fcn.ipynb b/notebooks/empirical_ntk_fcn.ipynb
@@ -80,7 +80,7 @@
       "source": [
         "# We need at least jaxlib-0.1.73 to avoid certain CUDA bugs when using `implementation=auto`\n",
         "!pip install -q --upgrade pip\n",
-        "!pip install -q jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
       ]
     },

diff --git a/notebooks/empirical_ntk_resnet.ipynb b/notebooks/empirical_ntk_resnet.ipynb
@@ -70,9 +70,8 @@
    },
    "outputs": [],
    "source": [
-    "# We need at least jaxlib-0.1.73 to avoid certain CUDA bugs when using `implementation=auto`\n",
     "!pip install -q --upgrade pip\n",
-    "!pip install -q jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+    "!pip install -q jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
     "!pip install -q git+https://github.com/google/flax\n",
     "!pip install -q git+https://www.github.com/google/neural-tangents"
    ]

diff --git a/notebooks/experimental/empirical_ntk_resnet_tf.ipynb b/notebooks/experimental/empirical_ntk_resnet_tf.ipynb
@@ -1,5 +1,15 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text",
+    "id": "view-in-github"
+   },
+   "source": [
+    "<a href=\"https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/experimental/empirical_ntk_resnet_tf.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -204,9 +214,8 @@
     }
    ],
    "source": [
-    "# We need at least jaxlib-0.1.73 to avoid certain CUDA bugs when using `implementation=auto`\n",
     "!pip install --upgrade pip\n",
-    "!pip install --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+    "!pip install --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
     "!pip install git+https://www.github.com/google/neural-tangents.git"
    ]
   },

diff --git a/notebooks/function_space_linearization.ipynb b/notebooks/function_space_linearization.ipynb
@@ -80,7 +80,7 @@
    "outputs": [],
    "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q tensorflow-datasets\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
    ]

diff --git a/notebooks/learning_curves.ipynb b/notebooks/learning_curves.ipynb
@@ -82,7 +82,7 @@
       ],
       "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
       ]
     },

diff --git a/notebooks/myrtle_kernel_with_neural_tangents.ipynb b/notebooks/myrtle_kernel_with_neural_tangents.ipynb
@@ -43,7 +43,7 @@
       "outputs": [],
       "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
       ]
     },

diff --git a/notebooks/neural_tangents_cookbook.ipynb b/notebooks/neural_tangents_cookbook.ipynb
@@ -53,7 +53,7 @@
    "outputs": [],
    "source": [
     "!pip install -q --upgrade pip\n",
-    "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+    "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
     "!pip install -q git+https://www.github.com/google/neural-tangents"
    ]
   },

diff --git a/notebooks/phase_diagram.ipynb b/notebooks/phase_diagram.ipynb
@@ -75,7 +75,7 @@
       ],
       "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
       ]
     },

diff --git a/notebooks/weight_space_linearization.ipynb b/notebooks/weight_space_linearization.ipynb
@@ -80,7 +80,7 @@
    "outputs": [],
    "source": [
         "!pip install -q --upgrade pip\n",
-        "!pip install -q --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install -q --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "!pip install -q tensorflow-datasets\n",
         "!pip install -q git+https://www.github.com/google/neural-tangents"
    ]