chex added (#330)

prishajain1 · web-flow · commit fcb1580b00b3 · 2026-02-13T08:50:20.000-08:00
diff --git a/requirements.txt b/requirements.txt
@@ -4,6 +4,7 @@ jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0
 absl-py
+chex
 datasets
 flax>=0.12.0
 optax>=0.2.3
diff --git a/requirements_with_jax_ai_image.txt b/requirements_with_jax_ai_image.txt
@@ -6,6 +6,7 @@ jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0
 absl-py
+chex
 datasets
 flax>=0.12.0
 optax>=0.2.3
