Add pytensor linker matrix ("c", "numba") to tests

ricardoV94 · ricardoV94 · commit 11964ce46126 · 2025-12-14T12:31:11.000+01:00
Remove useless os paramerizations and update codecov name
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -8,8 +8,6 @@ on:
     branches:
       - main
 
-
-
 # Cancel running workflows for updated PRs
 # https://turso.tech/blog/simple-trick-to-save-environment-and-money-when-using-github-actions
 concurrency:
@@ -26,7 +24,6 @@ concurrency:
 # enforces that test run just once per OS / floatX setting.
 
 jobs:
-
   changes:
     name: "Check for changes"
     runs-on: ubuntu-latest
@@ -56,8 +53,7 @@ jobs:
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        os: [ubuntu-latest]
-        floatx: [float64]
+        linker: [c, numba]
         python-version: ["3.13"]
         test-subset:
           - |
@@ -145,10 +141,10 @@ jobs:
             tests/dims/test_model.py
 
       fail-fast: false
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
+      PYTENSOR_FLAGS: linker=${{ matrix.linker }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -175,18 +171,17 @@ jobs:
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
         with:
-          token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
+          token: ${{ secrets.CODECOV_TOKEN }}
           env_vars: TEST_SUBSET
-          name: ${{ matrix.os }} ${{ matrix.floatx }}
+          name: Ubuntu py${{ matrix.python-version }} linker=${{ matrix.linker }}
           fail_ci_if_error: false
 
   windows:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        os: [windows-latest]
-        floatx: [float64]
+        linker: [c, numba]
         python-version: ["3.11"]
         test-subset:
           - tests/variational/test_approximations.py tests/variational/test_callbacks.py tests/variational/test_inference.py tests/variational/test_opvi.py tests/test_initial_point.py
@@ -195,10 +190,10 @@ jobs:
           - tests/step_methods/test_metropolis.py tests/step_methods/test_slicer.py tests/step_methods/hmc/test_nuts.py tests/step_methods/test_compound.py tests/step_methods/hmc/test_hmc.py tests/step_methods/test_state.py
 
       fail-fast: false
-    runs-on: ${{ matrix.os }}
+    runs-on: windows-latest
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
+      PYTENSOR_FLAGS: linker=${{ matrix.linker }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -225,18 +220,17 @@ jobs:
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
         with:
-          token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
+          token: ${{ secrets.CODECOV_TOKEN }}
           env_vars: TEST_SUBSET
-          name: ${{ matrix.os }} ${{ matrix.floatx }}
+          name: Windows py${{ python-version }} linker=${{ matrix.linker }}
           fail_ci_if_error: false
 
   macos:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        os: [macos-latest]
-        floatx: [float64]
+        linker: [c, numba]
         python-version: ["3.13"]
         test-subset:
           - |
@@ -253,10 +247,10 @@ jobs:
             tests/backends/test_zarr.py
             tests/variational/test_updates.py
       fail-fast: false
-    runs-on: ${{ matrix.os }}
+    runs-on: macos-latest
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
+      PYTENSOR_FLAGS: linker=${{ matrix.linker }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -285,16 +279,15 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
           env_vars: TEST_SUBSET
-          name: ${{ matrix.os }} ${{ matrix.floatx }}
+          name: MacOS py${{ python-version }} linker=${{ matrix.linker }}
           fail_ci_if_error: false
 
   alternative_backends:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        os: [ubuntu-latest]
-        floatx: [float64]
+        linker: [c, numba]
         python-version: ["3.13"]
         test-subset:
           - |
@@ -303,10 +296,10 @@ jobs:
             tests/sampling/test_mcmc_external.py
 
       fail-fast: false
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
+      PYTENSOR_FLAGS: linker=${{ matrix.linker }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -335,24 +328,24 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
           env_vars: TEST_SUBSET
-          name: Alternative backend tests - ${{ matrix.os }} ${{ matrix.floatx }}
+          name: Alternative backends py${{ python-version }} linker=${{ matrix.linker }}
           fail_ci_if_error: false
 
   float32:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
+        linker: [c, numba]
         os: [windows-latest]
-        floatx: [float32]
         python-version: ["3.13"]
         test-subset:
-        - tests/sampling/test_mcmc.py tests/ode/test_ode.py tests/ode/test_utils.py tests/distributions/test_transform.py
+          - tests/sampling/test_mcmc.py tests/ode/test_ode.py tests/ode/test_utils.py tests/distributions/test_transform.py
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
+      PYTENSOR_FLAGS: floatX=float32
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -381,19 +374,19 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
           env_vars: TEST_SUBSET
-          name: ${{ matrix.os }} ${{ matrix.floatx }}
+          name: float32 ${{ matrix.os }} py${{ matrix.python-version }} linker=${{ matrix.linker }}
           fail_ci_if_error: false
 
   all_tests:
     if: ${{ always() }}
     runs-on: ubuntu-latest
-    needs: [ changes, ubuntu, windows, macos, alternative_backends, float32 ]
+    needs: [changes, ubuntu, windows, macos, alternative_backends, float32]
     steps:
       - name: Check build matrix status
         if: ${{ needs.changes.outputs.changes == 'true' &&
-                ( needs.ubuntu.result != 'success' ||
-                  needs.windows.result != 'success' ||
-                  needs.macos.result != 'success' ||
-                  needs.alternative_backends.result != 'success' ||
-                  needs.float32.result != 'success' ) }}
+          ( needs.ubuntu.result != 'success' ||
+          needs.windows.result != 'success' ||
+          needs.macos.result != 'success' ||
+          needs.alternative_backends.result != 'success' ||
+          needs.float32.result != 'success' ) }}
         run: exit 1
diff --git a/scripts/check_all_tests_are_covered.py b/scripts/check_all_tests_are_covered.py
@@ -31,7 +31,7 @@ def find_testfiles():
 
 
 def from_yaml():
-    """Determine how often each test file is run per platform and floatX setting.
+    """Determine how often each test file is run per platform and linker setting.
 
     An exception is raised if tests run multiple times with the same configuration.
     """
@@ -41,48 +41,43 @@ def from_yaml():
         wfname = wf.rstrip(".yml")
         wfdef = yaml.safe_load(open(Path(".github", "workflows", wf)))
         for jobname, jobdef in wfdef["jobs"].items():
+            if jobname in ("float32", "all_tests"):
+                continue
+            runs_on = jobdef.get("runs-on", "unknown")
+            floatX = "float32" if jobname == "float32" else "float64"
             matrix = jobdef.get("strategy", {}).get("matrix", {})
             if matrix:
+                # Some jobs are parametrized by os, for others it's fixed
+                matrix.setdefault("os", [runs_on])
+                matrix.setdefault("floatX", [floatX])
                 matrices[(wfname, jobname)] = matrix
             else:
                 _log.warning("No matrix in %s/%s", wf, jobname)
 
-    # Now create an empty DataFrame to count based on OS/floatX/testfile
+    # Now create an empty DataFrame to count based on OS/linker/testfile
     all_os = []
-    all_floatX = []
+    all_linker = []
     for matrix in matrices.values():
         all_os += matrix["os"]
-        all_floatX += matrix["floatx"]
+        all_linker += matrix["linker"]
     all_os = tuple(sorted(set(all_os)))
-    all_floatX = tuple(sorted(set(all_floatX)))
+    all_linker = tuple(sorted(set(all_linker)))
     all_tests = find_testfiles()
 
     df = pandas.DataFrame(
         columns=pandas.MultiIndex.from_product(
-            [sorted(all_floatX), sorted(all_os)], names=["floatX", "os"]
+            [sorted(all_linker), sorted(all_os)], names=["linker", "os"]
         ),
         index=pandas.Index(sorted(all_tests), name="testfile"),
     )
     df.loc[:, :] = 0
 
     # Count how often the testfiles are included in job definitions
     for matrix in matrices.values():
-        for os_, floatX, subset in itertools.product(
-            matrix["os"], matrix["floatx"], matrix["test-subset"]
+        for os_, linker, subset in itertools.product(
+            matrix["os"], matrix["linker"], matrix["test-subset"]
         ):
             lines = [k for k in subset.split("\n") if k]
-            if "windows" in os_:
-                # Windows jobs need \ in line breaks within the test-subset!
-                # The following checks that these trailing \ are present in
-                # all items except the last.
-                if lines and lines[-1].endswith(" \\"):
-                    raise Exception(
-                        f"Last entry '{lines}' in Windows test subset should end WITHOUT ' \\'."
-                    )
-                for line in lines[:-1]:
-                    if not line.endswith(" \\"):
-                        raise Exception(f"Missing ' \\' after '{line}' in Windows test-subset.")
-                lines = [line.rstrip(" \\") for line in lines]
 
             # Unpack lines with >1 item
             testfiles = []
@@ -97,7 +92,7 @@ def from_yaml():
                 included = all_tests - ignored
 
             for testfile in included:
-                df.loc[testfile, (floatX, os_)] += 1
+                df.loc[testfile, (linker, os_)] += 1
 
     ignored_by_all = set(df[df.eq(0).all(axis=1)].index)
     run_multiple_times = set(df[df.gt(1).any(axis=1)].index)
@@ -111,7 +106,7 @@ def from_yaml():
         )
     if run_multiple_times:
         raise AssertionError(
-            f"{len(run_multiple_times)} tests are run multiple times with the same OS and floatX setting:\n{run_multiple_times}"
+            f"{len(run_multiple_times)} tests are run multiple times with the same OS and pytensor flags:\n{run_multiple_times}"
         )
     return
 
