[Bug Fix] Fix a nit issue for Jax rollback (#856)

* snapshot * snapshot
apple · Nov 25, 2024 · b352259 · b352259
1 parent d0edead
commit b352259
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -102,7 +102,7 @@ gcp = [
 # Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
 tpu = [
     "axlearn[gcp]",
-    "jax[tpu]==0.4.34",  # must be >=0.4.19 for compat with v5p.
+    "jax[tpu]==0.4.33",  # must be >=0.4.19 for compat with v5p.
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
 vertexai_tensorboard = [