From 53e572caa3144f77951e233f6bcca72b63084bd8 Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Fri, 24 Sep 2021 00:48:58 +0200
Subject: [PATCH] Fix CI by pinning upper JAX version number temporarily

Track https://github.com/pyro-ppl/numpyro/issues/1156 to see when we can remove the pin again.
---
 .github/workflows/jaxtests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/jaxtests.yml b/.github/workflows/jaxtests.yml
index 0e57bba3c6..5e7fc3d212 100644
--- a/.github/workflows/jaxtests.yml
+++ b/.github/workflows/jaxtests.yml
@@ -58,7 +58,7 @@ jobs:
       - name: Install jax specific dependencies
         run: |
           conda activate pymc3-dev-py39
-          pip install numpyro tensorflow_probability
+          pip install numpyro tensorflow_probability "jax<0.2.21"
       - name: Run tests
         run: |
           python -m pytest -vv --cov=pymc3 --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET