Merge pull request #12761 from skye:colab_tpu_driver

PiperOrigin-RevId: 480499024
jax-ml · Oct 12, 2022 · 012398b · 012398b
2 parents 6b459f5 + e2aa939
commit 012398b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,8 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
 -->
 
 ## jax 0.3.23
+* Changes
+  * Update Colab TPU driver version for new jaxlib release.
 
 ## jaxlib 0.3.23
 

diff --git a/jax/tools/colab_tpu.py b/jax/tools/colab_tpu.py
@@ -22,16 +22,15 @@
 TPU_DRIVER_MODE = 0
 
 
-def setup_tpu(tpu_driver_version='tpu_driver-0.2'):
+def setup_tpu(tpu_driver_version='tpu_driver_20221011'):
   """Sets up Colab to run on TPU.
 
   Note: make sure the Colab Runtime is set to Accelerator: TPU.
 
   Args
   ----
   tpu_driver_version : (str) specify the version identifier for the tpu driver.
-    Defaults to "tpu_driver-0.2", which can be used with jaxlib 0.3.20. Set to
-    "tpu_driver_nightly" to use the nightly tpu driver build.
+    Set to "tpu_driver_nightly" to use the nightly tpu driver build.
   """
   global TPU_DRIVER_MODE