Change default linker back to CVM

ricardoV94 · ricardoV94 · commit c9dd687ea801 · 2025-12-09T15:51:03.000+01:00
* Allow linker to change at runtime, and affect both `Mode` and `FAST_RUN` modes
* Remove float32 parametrization from test suite. Tests that care about this should cover it explicitly
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -65,7 +65,7 @@ jobs:
       - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
 
   test:
-    name: "${{ matrix.os }} test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}"
+    name: "tests : default-mode ${{ matrix.default-mode }} : py${{ matrix.python-version }} : ${{ matrix.os }} ${{ matrix.part }}"
     needs:
       - changes
       - style
@@ -74,11 +74,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ["ubuntu-latest"]
+        default-mode: [ "C", "NUMBA", "FAST_COMPILE" ]
         python-version: ["3.11", "3.14"]
-        fast-compile: [0, 1]
-        float32: [0, 1]
-        install-numba: [0]
+        os: [ "ubuntu-latest" ]
         install-jax: [0]
         install-torch: [0]
         install-mlx: [0]
@@ -94,81 +92,38 @@ jobs:
           - "tests/tensor/linalg tests/tensor/test_nlinalg.py tests/tensor/test_slinalg.py"
         exclude:
           - python-version: "3.11"
-            fast-compile: 1
-          - python-version: "3.11"
-            float32: 1
-          - fast-compile: 1
-            float32: 1
+            default-mode: "FAST_COMPILE"
         include:
-          - os: "ubuntu-latest"
-            part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link --ignore=pytensor/ipython.py"
+          - part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link --ignore=pytensor/ipython.py"
             python-version: "3.12"
-            fast-compile: 0
-            float32: 0
-            install-numba: 0
-            install-jax: 0
-            install-torch: 0
-            install-mlx: 0
-            install-xarray: 0
-          - install-numba: 1
-            os: "ubuntu-latest"
-            python-version: "3.11"
-            fast-compile: 0
-            float32: 0
-            part: "tests/link/numba --ignore=tests/link/numba/test_slinalg.py"
-          - install-numba: 1
-            os: "ubuntu-latest"
-            python-version: "3.14"
-            fast-compile: 0
-            float32: 0
-            part: "tests/link/numba --ignore=tests/link/numba/test_slinalg.py"
-          - install-numba: 1
-            os: "ubuntu-latest"
-            python-version: "3.14"
-            fast-compile: 0
-            float32: 0
-            part: "tests/link/numba/test_slinalg.py"
-          - install-jax: 1
-            os: "ubuntu-latest"
-            python-version: "3.11"
-            fast-compile: 0
-            float32: 0
-            part: "tests/link/jax"
-          - install-jax: 1
-            os: "ubuntu-latest"
+            default-mode: "C"
+          - part: "tests/link/numba --ignore=tests/link/numba/test_slinalg.py"
+            python-version: "3.12"
+            default-mode: "C"
+          - part: "tests/link/numba/test_slinalg.py"
+            python-version: "3.13"
+            default-mode: "C"
+          - part: "tests/link/jax"
+            install-jax: 1
             python-version: "3.14"
-            fast-compile: 0
-            float32: 0
-            part: "tests/link/jax"
-          - install-torch: 1
-            os: "ubuntu-latest"
+            default-mode: "C"
+          - part: "tests/link/pytorch"
+            install-torch: 1
             python-version: "3.11"
-            fast-compile: 0
-            float32: 0
-            part: "tests/link/pytorch"
-          - install-xarray: 1
-            os: "ubuntu-latest"
+            default-mode: "C"
+          - part: "tests/xtensor"
+            install-xarray: 1
             python-version: "3.14"
-            fast-compile: 0
-            float32: 0
-            part: "tests/xtensor"
-          - os: "macos-15"
-            python-version: "3.11"
-            fast-compile: 0
-            float32: 0
+            default-mode: "C"
+          - part: "tests/link/mlx"
             install-mlx: 1
-            install-numba: 0
-            install-jax: 0
-            install-torch: 0
-            part: "tests/link/mlx"
-          - os: "macos-15"
+            os: "macos-15"
+            python-version: "3.11"
+            default-mode: "C"
+          - part: "tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py tests/tensor/test_blas.py"
+            os: "macos-15"
             python-version: "3.14"
-            fast-compile: 0
-            float32: 0
-            install-numba: 0
-            install-jax: 0
-            install-torch: 0
-            part: "tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py tests/tensor/test_blas.py"
+            default-mode: "C"
 
     steps:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
@@ -199,11 +154,10 @@ jobs:
         run: |
 
           if [[ $OS == "macos-15" ]]; then
-            micromamba install --yes -q "python~=${PYTHON_VERSION}" "numpy${NUMPY_VERSION}" scipy pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate;
+            micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy scipy "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate;
           else
-            micromamba install --yes -q "python~=${PYTHON_VERSION}" mkl "numpy${NUMPY_VERSION}" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx;
+            micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy scipy "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service;
           fi
-          pip install "numba>=0.63"
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" jax jaxlib numpyro equinox && pip install tfp-nightly; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
           if [[ $INSTALL_MLX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "mlx<0.29.4"; fi
@@ -219,7 +173,6 @@ jobs:
           fi
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
-          INSTALL_NUMBA: ${{ matrix.install-numba }}
           INSTALL_JAX: ${{ matrix.install-jax }}
           INSTALL_TORCH: ${{ matrix.install-torch}}
           INSTALL_XARRAY: ${{ matrix.install-xarray }}
@@ -229,8 +182,8 @@ jobs:
       - name: Run tests
         shell: micromamba-shell {0}
         run: |
-          if [[ $FAST_COMPILE == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi
-          if [[ $FLOAT32 == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,floatX=float32; fi
+          if [[ $default-mode == "FAST_COMPILE" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi
+          if [[ $default-mode == "NUMBA" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,fast_run_backend=NUMBA; fi
           export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
           python -m pytest -r A --verbose --runslow --durations=50 --cov=pytensor/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART --benchmark-skip
         env:
@@ -239,8 +192,8 @@ jobs:
           MKL_NUM_THREADS: 1
           OMP_NUM_THREADS: 1
           PART: ${{ matrix.part }}
-          FAST_COMPILE: ${{ matrix.fast-compile }}
-          FLOAT32: ${{ matrix.float32 }}
+          default-mode: ${{ matrix.default-mode }}
+
 
       - name: Upload coverage file
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
diff --git a/pytensor/compile/__init__.py b/pytensor/compile/__init__.py
@@ -17,8 +17,8 @@
 )
 from pytensor.compile.io import In, Out, SymbolicInput, SymbolicOutput
 from pytensor.compile.mode import (
+    CVM,
     FAST_COMPILE,
-    FAST_RUN,
     JAX,
     NUMBA,
     OPT_FAST_COMPILE,
@@ -33,6 +33,7 @@
     PYTORCH,
     AddDestroyHandler,
     AddFeatureOptimizer,
+    C,
     Mode,
     PrintCurrentFunctionGraph,
     get_default_mode,
diff --git a/pytensor/compile/mode.py b/pytensor/compile/mode.py
@@ -5,7 +5,7 @@
 
 import logging
 import warnings
-from typing import Literal
+from typing import Any, Literal
 
 from pytensor.compile.function.types import Supervisor
 from pytensor.configdefaults import config
@@ -62,20 +62,17 @@ def register_linker(name, linker):
     predefined_linkers[name] = linker
 
 
-exclude = []
-if not config.cxx:
-    exclude = ["cxx_only"]
-OPT_NONE = RewriteDatabaseQuery(include=[], exclude=exclude)
+OPT_NONE = RewriteDatabaseQuery(include=[])
 # Minimum set of rewrites needed to evaluate a function. This is needed for graphs with "dummy" Operations
-OPT_MINIMUM = RewriteDatabaseQuery(include=["minimum_compile"], exclude=exclude)
+OPT_MINIMUM = RewriteDatabaseQuery(include=["minimum_compile"])
 # Even if multiple merge optimizer call will be there, this shouldn't
 # impact performance.
-OPT_MERGE = RewriteDatabaseQuery(include=["merge"], exclude=exclude)
-OPT_FAST_RUN = RewriteDatabaseQuery(include=["fast_run"], exclude=exclude)
+OPT_MERGE = RewriteDatabaseQuery(include=["merge"])
+OPT_FAST_RUN = RewriteDatabaseQuery(include=["fast_run"])
 OPT_FAST_RUN_STABLE = OPT_FAST_RUN.requiring("stable")
 
-OPT_FAST_COMPILE = RewriteDatabaseQuery(include=["fast_compile"], exclude=exclude)
-OPT_STABILIZE = RewriteDatabaseQuery(include=["fast_run"], exclude=exclude)
+OPT_FAST_COMPILE = RewriteDatabaseQuery(include=["fast_compile"])
+OPT_STABILIZE = RewriteDatabaseQuery(include=["fast_run"])
 OPT_STABILIZE.position_cutoff = 1.5000001
 OPT_NONE.name = "OPT_NONE"
 OPT_MINIMUM.name = "OPT_MINIMUM"
@@ -313,6 +310,8 @@ def __init__(
     ):
         if linker is None:
             linker = config.linker
+        if isinstance(linker, str) and linker == "auto":
+            linker = "cvm" if config.cxx else "vm"
         if isinstance(optimizer, str) and optimizer == "default":
             optimizer = config.optimizer
 
@@ -448,20 +447,15 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
         return new_mode
 
 
+C = Mode("c", "fast_run")
+CVM = Mode("cvm", "fast_run")
+VM = (Mode("vm", "fast_run"),)
+
 NUMBA = Mode(
     NumbaLinker(),
     RewriteDatabaseQuery(include=["fast_run", "numba"]),
 )
 
-FAST_COMPILE = Mode(
-    NumbaLinker(),
-    RewriteDatabaseQuery(include=["fast_compile"]),
-)
-FAST_RUN = NUMBA
-
-C = Mode("c", "fast_run")
-CVM = Mode("cvm", "fast_run")
-
 JAX = Mode(
     JAXLinker(),
     RewriteDatabaseQuery(include=["fast_run", "jax"]),
@@ -476,10 +470,19 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     RewriteDatabaseQuery(include=["fast_run"]),
 )
 
+FAST_COMPILE = Mode(
+    VMLinker(use_cloop=False, c_thunks=False),
+    RewriteDatabaseQuery(include=["fast_compile", "py_only"]),
+)
+
+fast_run_linkers_to_mode = {
+    "cvm": CVM,
+    "vm": VM,
+    "numba": NUMBA,
+}
 
 predefined_modes = {
     "FAST_COMPILE": FAST_COMPILE,
-    "FAST_RUN": FAST_RUN,
     "C": C,
     "CVM": CVM,
     "JAX": JAX,
@@ -488,7 +491,7 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     "MLX": MLX,
 }
 
-_CACHED_RUNTIME_MODES: dict[str, Mode] = {}
+_CACHED_RUNTIME_MODES: dict[Any, Mode] = {}
 
 
 def get_mode(orig_string):
@@ -506,10 +509,20 @@ def get_mode(orig_string):
     if upper_string in predefined_modes:
         return predefined_modes[upper_string]
 
+    if upper_string == "FAST_RUN":
+        linker = config.linker
+        if linker == "auto":
+            return CVM if config.cxx else VM
+        return fast_run_linkers_to_mode[linker]
+
     global _CACHED_RUNTIME_MODES
 
-    if upper_string in _CACHED_RUNTIME_MODES:
-        return _CACHED_RUNTIME_MODES[upper_string]
+    cache_key = ("MODE", config.linker) if upper_string == "MODE" else upper_string
+
+    try:
+        return _CACHED_RUNTIME_MODES[cache_key]
+    except KeyError:
+        pass
 
     # Need to define the mode for the first time
     if upper_string == "MODE":
@@ -535,7 +548,7 @@ def get_mode(orig_string):
     if config.optimizer_requiring:
         ret = ret.requiring(*config.optimizer_requiring.split(":"))
     # Cache the mode for next time
-    _CACHED_RUNTIME_MODES[upper_string] = ret
+    _CACHED_RUNTIME_MODES[cache_key] = ret
 
     return ret
 
diff --git a/pytensor/configdefaults.py b/pytensor/configdefaults.py
@@ -371,24 +371,26 @@ def add_compile_configvars():
     )
     del param
 
-    default_linker = "numba"
+    default_linker = "auto"
 
     if rc == 0 and config.cxx != "":
         # Keep the default linker the same as the one for the mode FAST_RUN
         linker_options = [
-            "cvmc|py",
+            "cvm",
+            "c|py",
             "py",
             "c",
             "c|py_nogc",
             "vm",
             "vm_nogc",
             "cvm_nogc",
+            "numba",
             "jax",
         ]
     else:
         # g++ is not present or the user disabled it,
         # linker should default to python only.
-        linker_options = ["py", "vm", "vm_nogc", "jax"]
+        linker_options = ["py", "vm", "vm_nogc", "numba", "jax"]
         if type(config).cxx.is_default:
             # If the user provided an empty value for cxx, do not warn.
             _logger.warning(
@@ -400,9 +402,8 @@ def add_compile_configvars():
 
     config.add(
         "linker",
-        "Default linker used if the pytensor flags mode is Mode",
-        # Not mutable because the default mode is cached after the first use.
-        EnumStr(default_linker, linker_options, mutable=False),
+        "Default linker used if the pytensor flags mode is Mode or FAST_RUN",
+        EnumStr(default_linker, linker_options, mutable=True),
         in_c_key=False,
     )
 
diff --git a/pytensor/configparser.py b/pytensor/configparser.py
@@ -76,6 +76,7 @@ class PyTensorConfigParser:
     unpickle_function: bool
     # add_compile_configvars
     mode: str
+    fast_run_backend: str
     cxx: str
     linker: str
     allow_gc: bool
diff --git a/tests/compile/test_mode.py b/tests/compile/test_mode.py
@@ -77,9 +77,18 @@ def test_modes(self):
 
         # Linkers to use with regular Mode
         if config.cxx:
-            linkers = ["py", "c|py", "c|py_nogc", "vm", "vm_nogc", "cvm", "cvm_nogc"]
+            linkers = [
+                "py",
+                "c|py",
+                "c|py_nogc",
+                "vm",
+                "vm_nogc",
+                "cvm",
+                "cvm_nogc",
+                "numba",
+            ]
         else:
-            linkers = ["py", "c|py", "c|py_nogc", "vm", "vm_nogc"]
+            linkers = ["py", "c|py", "c|py_nogc", "vm", "vm_nogc", "numba"]
         modes = predef_modes + [Mode(linker, "fast_run") for linker in linkers]
 
         for mode in modes:
