Try to run full test suite in Numba backend

Author: ricardoV94

.github/workflows/test.yml

@@ -202,7 +202,7 @@ jobs:
           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;
           fi
-          if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "numba>=0.57"; fi
+          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "numba>=0.57";
           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

pytensor/compile/mode.py

@@ -65,9 +65,7 @@ def register_linker(name, linker):
 # If a string is passed as the optimizer argument in the constructor
 # for Mode, it will be used as the key to retrieve the real optimizer
 # in this dictionary
-exclude = []
-if not config.cxx:
-    exclude = ["cxx_only"]
+exclude = ["cxx_only", "BlasOpt"]
 OPT_NONE = RewriteDatabaseQuery(include=[], exclude=exclude)
 # 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)
@@ -445,37 +443,37 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
         return new_mode
 
 
-# If a string is passed as the mode argument in function or
-# FunctionMaker, the Mode will be taken from this dictionary using the
-# string as the key
-# Use VM_linker to allow lazy evaluation by default.
+numba_exclude = [
+    "cxx_only",
+    "BlasOpt",
+    "local_careduce_fusion",
+    "scan_save_mem_prealloc",
+]
+
 FAST_COMPILE = Mode(
-    VMLinker(use_cloop=False, c_thunks=False),
-    RewriteDatabaseQuery(include=["fast_compile", "py_only"]),
+    NumbaLinker(),
+    # TODO: Fast_compile should just use python code, CHANGE ME!
+    RewriteDatabaseQuery(
+        include=["fast_compile", "numba"],
+        exclude=numba_exclude,
+    ),
+)
+FAST_RUN = Mode(
+    NumbaLinker(),
+    RewriteDatabaseQuery(include=["fast_run", "numba"], exclude=numba_exclude),
 )
-if config.cxx:
-    FAST_RUN = Mode("cvm", "fast_run")
-else:
-    FAST_RUN = Mode(
-        "vm",
-        RewriteDatabaseQuery(include=["fast_run", "py_only"]),
-    )
-
-C = Mode("c", "fast_run")
-C_VM = Mode("cvm", "fast_run")
 
 NUMBA = Mode(
     NumbaLinker(),
     RewriteDatabaseQuery(
         include=["fast_run", "numba"],
-        exclude=[
-            "cxx_only",
-            "BlasOpt",
-            "local_careduce_fusion",
-            "scan_save_mem_prealloc",
-        ],
+        exclude=numba_exclude,
     ),
 )
+del numba_exclude
+
+C = Mode("c", "fast_run")
+C_VM = Mode("cvm", "fast_run")
 
 JAX = Mode(
     JAXLinker(),
@@ -523,7 +521,9 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     ),
 )
 
-
+# If a string is passed as the mode argument in function or
+# FunctionMaker, the Mode will be taken from this dictionary using the
+# string as the key
 predefined_modes = {
     "FAST_COMPILE": FAST_COMPILE,
     "FAST_RUN": FAST_RUN,

pytensor/configdefaults.py

@@ -371,25 +371,24 @@ def add_compile_configvars():
     )
     del param
 
-    default_linker = "cvm"
+    default_linker = "numba"
 
     if rc == 0 and config.cxx != "":
         # Keep the default linker the same as the one for the mode FAST_RUN
         linker_options = [
-            "c|py",
+            "cvmc|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", "numba", "jax"]
+        linker_options = ["py", "vm", "vm_nogc", "jax"]
         if type(config).cxx.is_default:
             # If the user provided an empty value for cxx, do not warn.
             _logger.warning(