diff --git a/requirements.txt b/requirements.txt
index 0516b9f2..888f6a7a 100644
--- 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
index c279edb8..d8776971 100644
--- 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