From 885e1ba8505b3607292c145b2b0b2b4ee1895c38 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Wed, 29 Mar 2023 15:18:26 +0200
Subject: [PATCH 1/2] Pin `chex` version

---
 .github/workflows/tests.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 405b32978..e8371883e 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -55,11 +55,12 @@ jobs:
 
     - name: Install dependencies
       # `jax[cuda]<0.4` because of Docker issues: https://github.com/google/jax/issues/13758
+      # `chex<0.1.7` because it requires `jax>=0.4.6`
       # `flax<0.6.5` because it requires `jax>=0.4.2`
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install -e".[test]"
-        python3 -m pip install "flax<0.6.5"
+        python3 -m pip install "flax<0.6.5 chex<0.1.7"
         python3 -m pip install "jax[cuda]<0.4" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
     - name: Nvidia SMI

From 31b8fcb4d63dfc6c92c211977fa7b3c968c6f656 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Wed, 29 Mar 2023 15:24:42 +0200
Subject: [PATCH 2/2] Fix typo

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

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index e8371883e..b139d48df 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -60,7 +60,7 @@ jobs:
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install -e".[test]"
-        python3 -m pip install "flax<0.6.5 chex<0.1.7"
+        python3 -m pip install "flax<0.6.5" "chex<0.1.7"
         python3 -m pip install "jax[cuda]<0.4" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
     - name: Nvidia SMI