Downstream 1285: Add a fusion rewrite for CAReduces with Elemwise inputs

pytensor/compile/function/pfunc.py

@@ -3,7 +3,6 @@
 
 """
 
-import logging
 from copy import copy
 from typing import Optional
 
@@ -16,11 +15,6 @@
 from pytensor.graph.fg import FunctionGraph
 
 
-_logger = logging.getLogger("pytensor.compile.function.pfunc")
-
-__docformat__ = "restructuredtext en"
-
-
 def rebuild_collect_shared(
     outputs,
     inputs=None,

pytensor/compile/function/types.py

@@ -32,6 +32,7 @@
 
 
 if TYPE_CHECKING:
+    from pytensor.compile.mode import Mode
     from pytensor.link.vm import VM
 
 
@@ -1391,16 +1392,24 @@ def check_unused_inputs(inputs, outputs, on_unused_input):
 
     @staticmethod
     def prepare_fgraph(
-        inputs, outputs, additional_outputs, fgraph, rewriter, linker, profile
+        inputs,
+        outputs,
+        additional_outputs,
+        fgraph: FunctionGraph,
+        mode: "Mode",
+        profile,
     ):
 
+        rewriter = mode.optimizer
+
         try:
             start_rewriter = time.perf_counter()
 
             rewriter_profile = None
             rewrite_time = None
 
             with config.change_flags(
+                mode=mode,
                 compute_test_value=config.compute_test_value_opt,
                 traceback__limit=config.traceback__compile_limit,
             ):
@@ -1440,7 +1449,7 @@ def prepare_fgraph(
                     stacklevel=3,
                 )
 
-        if not hasattr(linker, "accept"):
+        if not hasattr(mode.linker, "accept"):
             raise ValueError(
                 "'linker' parameter of FunctionMaker should be "
                 f"a Linker with an accept method or one of {list(pytensor.compile.mode.predefined_linkers.keys())}"
@@ -1511,12 +1520,8 @@ def __init__(
 
         self.fgraph = fgraph
 
-        rewriter, linker = mode.optimizer, copy.copy(mode.linker)
-
         if not no_fgraph_prep:
-            self.prepare_fgraph(
-                inputs, outputs, found_updates, fgraph, rewriter, linker, profile
-            )
+            self.prepare_fgraph(inputs, outputs, found_updates, fgraph, mode, profile)
 
         assert len(fgraph.outputs) == len(outputs + found_updates)
 
@@ -1528,6 +1533,8 @@ def __init__(
             if not spec.borrow
         ]
 
+        linker = copy.copy(mode.linker)
+
         if no_borrow:
             self.linker = linker.accept(
                 fgraph,

pytensor/compile/mode.py

@@ -7,6 +7,8 @@
 import warnings
 from typing import Optional, Tuple, Union
 
+from typing_extensions import Literal
+
 from pytensor.compile.function.types import Supervisor
 from pytensor.configdefaults import config
 from pytensor.graph.destroyhandler import DestroyHandler
@@ -530,3 +532,26 @@ def register_mode(name, mode):
     if name in predefined_modes:
         raise ValueError(f"Mode name already taken: {name}")
     predefined_modes[name] = mode
+
+
+def get_target_language(mode=None) -> Tuple[Literal["py", "c", "numba", "jax"], ...]:
+    """Get the compilation target language."""
+
+    if mode is None:
+        mode = get_default_mode()
+
+    linker = mode.linker
+
+    if isinstance(linker, NumbaLinker):
+        return ("numba",)
+    if isinstance(linker, JAXLinker):
+        return ("jax",)
+    if isinstance(linker, PerformLinker):
+        return ("py",)
+    if isinstance(linker, CLinker):
+        return ("c",)
+
+    if isinstance(linker, (VMLinker, OpWiseCLinker)):
+        return ("c", "py") if config.cxx else ("py",)
+
+    raise Exception(f"Unsupported Linker: {linker}")