From 2fecf870eba13bb1eee08cd44333ba2674f4b0e0 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Wed, 15 Jul 2026 22:37:31 -0700 Subject: [PATCH] fix: pin jax/jaxlib/flax lower bounds to avoid linear_util ImportError pyproject.toml and requirements.txt pinned jax, jaxlib and flax with no lower bounds, so the resolver could select an old Flax (<0.8) that does 'from jax import linear_util' against a modern JAX. jax.linear_util was removed in JAX 0.4.24, producing 'ImportError: cannot import name linear_util' at import time. Add floors (jax>=0.4.30, jaxlib>=0.4.30, flax>=0.8.0) so a compatible set is always resolved. Closes #24 Signed-off-by: ayaangazali --- pyproject.toml | 6 +++--- requirements.txt | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 08b0472..d6e57c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,9 +3,9 @@ name = "needle" version = "0.1.0" requires-python = ">=3.11" dependencies = [ - "jax", - "jaxlib", - "flax", + "jax>=0.4.30", + "jaxlib>=0.4.30", + "flax>=0.8.0", "optax", "datasets", "huggingface_hub", diff --git a/requirements.txt b/requirements.txt index eaf2d02..c361064 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ -jax -jaxlib -flax +jax>=0.4.30 +jaxlib>=0.4.30 +flax>=0.8.0 optax datasets huggingface_hub