Bump jax version to 0.8.0 (#1006)

py4 · Pooya Moradi · web-flow · commit 92f94be3920b · 2025-11-04T15:01:34.000-08:00
Signed-off-by: Pooya Moradi &lt;pooyam@google.com&gt;
Co-authored-by: Pooya Moradi &lt;pooyam@google.com&gt;
diff --git a/requirements.txt b/requirements.txt
@@ -5,9 +5,8 @@ pytest-mock
 absl-py
 numpy
 google-cloud-storage
-jax==0.7.2
-jaxlib==0.7.2
-libtpu==0.0.23
+jax[tpu]==0.8.0
+jaxlib==0.8.0
 jaxtyping
 flax==0.11.1
 torchax==0.0.7
