Upgrade to numpyro version with bugfixes

See #5020

Author: michaelosthege

.github/workflows/jaxtests.yml

@@ -58,7 +58,7 @@ jobs:
       - name: Install jax specific dependencies
         run: |
           conda activate pymc-dev-py39
-          pip install numpyro tensorflow_probability "jax<0.2.21"
+          pip install "numpyro>=0.8.0" tensorflow_probability
       - name: Run tests
         run: |
           python -m pytest -vv --cov=pymc --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET