change jax version

susanbao · susanbao · commit 01712be5e7f1 · 2025-10-08T23:45:17.000Z
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
-jax>=0.6.2
+jax==0.7.0
 jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0
