Run numba and jax tests in separate jobs

Author: ricardoV94

.github/workflows/test.yml

@@ -73,9 +73,10 @@ jobs:
         python-version: ["3.8", "3.11"]
         fast-compile: [0,1]
         float32: [0,1]
-        install-numba: [1]
+        install-numba: [0]
+        install-jax: [0]
         part:
-          - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse"
+          - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse --ignore=tests/link/numba --ignore=tests/link/jax"
           - "tests/scan"
           - "tests/sparse"
           - "tests/tensor --ignore=tests/tensor/conv --ignore=tests/tensor/rewriting --ignore=tests/tensor/test_math.py --ignore=tests/tensor/test_basic.py --ignore=tests/tensor/test_blas.py --ignore=tests/tensor/test_math_scipy.py --ignore=tests/tensor/test_inplace.py --ignore=tests/tensor/test_elemwise.py"
@@ -93,6 +94,15 @@ jobs:
             part: "tests/tensor/test_math.py"
           - fast-compile: 1
             float32: 1
+        include:
+          - install-numba: 1
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/numba"
+          - install-jax: 1
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/jax"
     steps:
       - uses: actions/checkout@v3
         with:
@@ -118,8 +128,8 @@ jobs:
         shell: bash -l {0}
         run: |
           mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark sympy
-          if [[ $INSTALL_NUMBA == "1" ]]; then mamba install --yes -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numpy<1.23" "numba>=0.57" numba-scipy; fi
-          mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro
+          if [[ $INSTALL_NUMBA == "1" ]]; then mamba install --yes -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numpy<1.23" "numba>=0.57" numba-scipy "numpy=1.23"; fi
+          if [[ $INSTALL_JAX == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro; fi
           pip install -e ./
           mamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'