Split jax tests into their own workflow

.github/workflows/jaxtests.yml

@@ -0,0 +1,64 @@
+name: jax-sampling
+
+on:
+  pull_request:
+  push:
+    branches: [master]
+
+jobs:
+  pytest:
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        floatx: [float64]
+        test-subset:
+          - pymc3/tests/test_sampling_jax.py
+      fail-fast: false
+    runs-on: ${{ matrix.os }}
+    env:
+      TEST_SUBSET: ${{ matrix.test-subset }}
+      THEANO_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
+    defaults:
+      run:
+        shell: bash -l {0}
+    steps:
+      - uses: actions/checkout@v2
+      - name: Cache conda
+        uses: actions/cache@v1
+        env:
+          # Increase this value to reset cache if environment-dev-py39.yml has not changed
+          CACHE_NUMBER: 0
+        with:
+          path: ~/conda_pkgs_dir
+          key: ${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-${{
+            hashFiles('conda-envs/environment-dev-py39.yml') }}
+      - name: Cache multiple paths
+        uses: actions/cache@v2
+        env:
+          # Increase this value to reset cache if requirements.txt has not changed
+          CACHE_NUMBER: 0
+        with:
+          path: |
+            ~/.cache/pip
+            $RUNNER_TOOL_CACHE/Python/*
+            ~\AppData\Local\pip\Cache
+          key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{
+            hashFiles('requirements.txt') }}
+      - uses: conda-incubator/setup-miniconda@v2
+        with:
+          activate-environment: pymc3-dev-py39
+          channel-priority: strict
+          environment-file: conda-envs/environment-dev-py39.yml
+          use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
+      - name: Install pymc3
+        run: |
+          conda activate pymc3-dev-py39
+          pip install -e .
+          python --version
+      - name: Install jax specific dependencies
+        run: |
+          conda activate pymc3-dev-py39
+          pip install numpyro tensorflow_probability
+      - name: Run tests
+        run: |
+          python -m pytest -vv --cov=pymc3 --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET

.github/workflows/pytest.yml

@@ -27,6 +27,7 @@ jobs:
             --ignore=pymc3/tests/test_quadpotential.py
             --ignore=pymc3/tests/test_random.py
             --ignore=pymc3/tests/test_sampling.py
+            --ignore=pymc3/tests/test_sampling_jax.py
             --ignore=pymc3/tests/test_shape_handling.py
             --ignore=pymc3/tests/test_shared.py
             --ignore=pymc3/tests/test_smc.py