Drop latest ArviZ tests, upgrade to Ubuntu 20.04 and add MacOS jobs

Author: michaelosthege

.github/workflows/tests.yml

@@ -32,7 +32,7 @@ jobs:
   ubuntu:
     strategy:
       matrix:
-        os: [ubuntu-18.04]
+        os: [ubuntu-20.04]
         floatx: [float32, float64]
         test-subset:
         # How this works:
@@ -248,13 +248,24 @@ jobs:
           env_vars: TEST_SUBSET
           name: ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false
-  jax:
+  macos:
     strategy:
       matrix:
-        os: [ubuntu-latest]
+        os: [macos-latest]
         floatx: [float64]
         test-subset:
-          - pymc/tests/test_sampling_jax.py
+          - |
+            pymc/tests/test_parallel_sampling.py
+            pymc/tests/test_posteriors.py
+            pymc/tests/test_sampling.py
+
+          - |
+            pymc/tests/test_data_container.py
+            pymc/tests/test_idata_conversion.py
+            pymc/tests/test_missing.py
+            pymc/tests/test_model.py
+            pymc/tests/test_shape_handling.py
+            pymc/tests/test_shared.py
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
@@ -301,38 +312,22 @@ jobs:
           conda activate pymc-test-py39
           pip install -e .
           python --version
-      - name: Install jax specific dependencies
-        run: |
-          conda activate pymc-test-py39
-          pip install "numpyro>=0.8.0"
-          pip install git+https://github.com/blackjax-devs/blackjax.git@main
       - name: Run tests
         run: |
           python -m pytest -vv --cov=pymc --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
         with:
           env_vars: TEST_SUBSET
-          name: JAX tests - ${{ matrix.os }} ${{ matrix.floatx }}
+          name: ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false
-  arviz:
+  jax:
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-latest]
+        os: [ubuntu-20.04]
         floatx: [float64]
         test-subset:
-          - |
-            pymc/tests/test_parallel_sampling.py
-            pymc/tests/test_posteriors.py
-            pymc/tests/test_sampling.py
-
-          - |
-            pymc/tests/test_data_container.py
-            pymc/tests/test_idata_conversion.py
-            pymc/tests/test_missing.py
-            pymc/tests/test_model.py
-            pymc/tests/test_shape_handling.py
-            pymc/tests/test_shared.py
+          - pymc/tests/test_sampling_jax.py
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
@@ -379,17 +374,17 @@ jobs:
           conda activate pymc-test-py39
           pip install -e .
           python --version
-      - name: Install latest arviz
+      - name: Install jax specific dependencies
         run: |
           conda activate pymc-test-py39
-          pip uninstall arviz -y
-          pip install git+https://github.com/arviz-devs/arviz
+          pip install "numpyro>=0.8.0"
+          pip install git+https://github.com/blackjax-devs/blackjax.git@main
       - name: Run tests
         run: |
           python -m pytest -vv --cov=pymc --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
         with:
           env_vars: TEST_SUBSET
-          name: ArviZ compat - ${{ matrix.os }} ${{ matrix.floatx }}
+          name: JAX tests - ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false