Fix CI by pinning upper JAX version number temporarily

Track pyro-ppl/numpyro#1156 to see when we can remove the pin again.

Author: michaelosthege

.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