Set the global mode during compilation

Author: brandonwillard

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,