From 1be4b81e09b735ee2a9f68e8c3d4ebef31c2dac7 Mon Sep 17 00:00:00 2001 From: Prisha Jain Date: Tue, 10 Feb 2026 14:38:18 +0530 Subject: [PATCH] chex added --- requirements.txt | 1 + requirements_with_jax_ai_image.txt | 1 + 2 files changed, 2 insertions(+) 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