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

* snapshot * snapshot
qdavid1 · Dec 11, 2024 · 3e3d4d6 · 3e3d4d6
1 parent ff05661
commit 3e3d4d6
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 = [