higher flax version

allenai · Oct 2, 2023 · 49b2fc5 · 49b2fc5
1 parent 0b677cc
commit 49b2fc5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -211,7 +211,7 @@ jobs:
         if: steps.virtualenv-cache.outputs.cache-hit != 'true' && (contains(matrix.task.extras, 'flax') || contains(matrix.task.extras, 'all'))
         run: |
           . .venv/bin/activate
-          pip install flax==0.5.0 jax==0.4.1 jaxlib==0.4.1 tensorflow-cpu==2.9.1 optax==0.1.3
+          pip install flax==0.6.1 jax==0.4.1 jaxlib==0.4.1 tensorflow-cpu==2.9.1 optax==0.1.3
 
       - name: Install editable (no cache hit)
         if: steps.virtualenv-cache.outputs.cache-hit != 'true'

diff --git a/pyproject.toml b/pyproject.toml
@@ -90,7 +90,7 @@ fairscale = [
 flax = [
   "datasets>=1.12,<3.0",
   "jax>=0.4.1,<=0.4.13",
-  "flax>=0.5.0,<=0.7.0",
+  "flax>=0.6.1,<=0.7.0",
   "optax>=0.1.2",
   "tensorflow-cpu>=2.9.1"
 ]