jax, flax versions

allenai · Oct 2, 2023 · bea40bf · bea40bf
1 parent 64d61ac
commit bea40bf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -89,8 +89,8 @@ fairscale = [
 ]
 flax = [
   "datasets>=1.12,<3.0",
-  "jax>=0.3.13",
-  "flax>=0.5.0",
+  "jax>=0.4.1,<=0.4.13",
+  "flax>=0.5.0,<=0.7.0",
   "optax>=0.1.2",
   "tensorflow-cpu>=2.9.1"
 ]