Run numba and jax tests in separate jobs

Author: ricardoV94

.github/workflows/test.yml

@@ -73,7 +73,8 @@ 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/scan"
@@ -93,6 +94,27 @@ jobs:
             part: "tests/tensor/test_math.py"
           - fast-compile: 1
             float32: 1
+        include:
+          - install-numba: 1
+            python-version: "3.8"
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/numba"
+          - install-numba: 1
+            python-version: "3.11"
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/numba"
+          - install-jax: 1
+            python-version: "3.8"
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/jax"
+          - install-jax: 1
+            python-version: "3.11"
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/jax"
     steps:
       - uses: actions/checkout@v3
         with:
@@ -120,14 +142,15 @@ jobs:
           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" ]] && [[ $PYTHON_VERSION == "3.8" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numpy<1.23" "numba>=0.57" numba-scipy; fi
           if [[ $INSTALL_NUMBA == "1" ]] && [[ $PYTHON_VERSION != "3.8" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57" numba-scipy; fi
-          mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro
+          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))'
           python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
           INSTALL_NUMBA: ${{ matrix.install-numba }}
+          INSTALL_JAX: ${{ matrix.install-jax }}
 
       - name: Run tests
         shell: bash -l {0}

setup.cfg

@@ -9,6 +9,8 @@ per-file-ignores =
     pytensor/link/jax/jax_dispatch.py:E402,F403,F401
     pytensor/link/jax/jax_linker.py:E402,F403,F401
     pytensor/sparse/sandbox/sp2.py:F401
+    tests/link/jax/*.py:E402
+    tests/link/numba/*.py:E402
     tests/tensor/test_math_scipy.py:E402
     tests/sparse/test_basic.py:E402
     tests/sparse/test_opt.py:E402

tests/link/jax/test_tensor_basic.py

@@ -1,7 +1,10 @@
-import jax.errors
 import numpy as np
 import pytest
 
+
+jax = pytest.importorskip("jax")
+import jax.errors
+
 import pytensor
 import pytensor.tensor.basic as at
 from pytensor.configdefaults import config

tests/link/numba/test_basic.py

@@ -3,10 +3,12 @@
 from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, Tuple, Union
 from unittest import mock
 
-import numba
 import numpy as np
 import pytest
 
+
+numba = pytest.importorskip("numba")
+
 import pytensor.scalar as aes
 import pytensor.scalar.math as aesm
 import pytensor.tensor as at

tests/link/numba/test_cython_support.py

@@ -1,6 +1,11 @@
 import numpy as np
 import pytest
 import scipy.special.cython_special
+
+
+numba = pytest.importorskip("numba")
+
+
 from numba.types import float32, float64, int32, int64
 
 from pytensor.link.numba.dispatch.cython_support import Signature, wrap_cython_function

tests/link/numba/test_performance.py

@@ -3,6 +3,9 @@
 import numpy as np
 import pytest
 
+
+pytest.importorskip("numba")
+
 import pytensor.tensor as aet
 from pytensor import config
 from pytensor.compile.function import function
@@ -70,4 +73,5 @@ def test_careduce_performance(careduce_fn, numpy_fn, axis, inputs, input_vals):
     mean_numpy_time = np.mean(numpy_times)
     # mean_c_time = np.mean(c_times)
 
+    # FIXME: Why are we asserting >=? Numba could be doing worse than numpy!
     assert mean_numba_time / mean_numpy_time >= 0.75

tests/link/numba/test_sparse.py

@@ -1,8 +1,11 @@
-import numba
 import numpy as np
 import pytest
 import scipy as sp
 
+
+numba = pytest.importorskip("numba")
+
+
 # Make sure the Numba customizations are loaded
 import pytensor.link.numba.dispatch.sparse  # noqa: F401
 from pytensor import config