From f36919b10e2a4d74c4b7cb8ecafc19c9e8632484 Mon Sep 17 00:00:00 2001
From: Dirk Groeneveld <dirkg@allenai.org>
Date: Wed, 3 Apr 2024 15:37:11 -0700
Subject: [PATCH] Update jax dependencies in GitHub actions as well

---
 .github/workflows/main.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 5789133d5..3fe5e99cc 100644
--- 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.7" "jax<0.5" "jaxlib<0.5" "tensorflow-cpu<2.10" "optax<0.2"
+          pip install flax jax jaxlib "tensorflow-cpu>=2.9.1" optax
 
       - name: Install editable (no cache hit)
         if: steps.virtualenv-cache.outputs.cache-hit != 'true'