Pin JAX version until next release of jax_tpu_embedding (#106)

jeffcarp · web-flow · commit 24ba6a1a64ff · 2025-08-14T10:07:29.000-07:00
This should fix CI.
diff --git a/requirements.txt b/requirements.txt
@@ -1,11 +1,11 @@
 absl-py
 clu
-jax[cpu]
+jax[cpu]==0.6.2
 jax_tpu_embedding
 keras-hub
 keras-rs
 pytest
 rouge-score
 scikit-learn
 tensorflow
-torchmetrics
+torchmetrics
