Fix Minimize Ops type strictness

pytensor/compile/maker.py

@@ -469,14 +469,7 @@ def prepare_fgraph(
                 mode=mode,
                 traceback__limit=config.traceback__compile_limit,
             ):
-                # Expose the compile mode so inner-graph rewrites can recover
-                # the active linker's required/incompatible rewrites reliably
-                # (``config.mode`` is unreliable across nested compilations).
-                fgraph._compile_mode = mode
-                try:
-                    rewriter_profile = rewriter(fgraph)
-                finally:
-                    del fgraph._compile_mode
+                rewriter_profile = rewriter(fgraph)
 
                 end_rewriter = time.perf_counter()
                 rewrite_time = end_rewriter - start_rewriter

pytensor/compile/mode.py

@@ -4,6 +4,7 @@
 """
 
 import logging
+import sys
 import warnings
 from typing import Any
 
@@ -279,6 +280,53 @@ def apply(self, fgraph):
 del _tags
 
 
+class _ActiveModeGraphRewriter(GraphRewriter):
+    """Wrap a mode's optimizer so applying it records the active compile mode.
+
+    Exposes the compile mode as ``fgraph._compile_mode`` (see ``get_active_mode``)
+    so rewrites can make decisions based on the compile target (fusion split rules,
+    inner graph specialization, ...). Stamps the mode when the fgraph does not already carry one, and
+    restores the previous state after.
+    """
+
+    def __init__(self, rewriter: GraphRewriter, mode: "Mode"):
+        self._rewriter = rewriter
+        self._mode = mode
+
+    def add_requirements(self, fgraph):
+        return self._rewriter.add_requirements(fgraph)
+
+    def apply(self, fgraph, *args, **kwargs):
+        # Only stamp when unset, so a mode already established by an outer
+        # optimizer run (the outermost compilation, or a nested inner-graph
+        # rewrite) wins.
+        stamped = getattr(fgraph, "_compile_mode", None) is None
+        if stamped:
+            fgraph._compile_mode = self._mode
+        try:
+            return self._rewriter.apply(fgraph, *args, **kwargs)
+        finally:
+            if stamped:
+                del fgraph._compile_mode
+
+    def print_summary(self, stream=sys.stdout, level=0, depth=-1):
+        return self._rewriter.print_summary(stream=stream, level=level, depth=depth)
+
+    def print_profile(self, stream, prof, level=0):
+        # ``GraphRewriter.print_profile`` is a classmethod that raises unless
+        # overridden, so delegate to the wrapped rewriter's implementation.
+        return self._rewriter.print_profile(stream, prof, level=level)
+
+    def __getattr__(self, name):
+        # Delegate everything else (profiling helpers, query metadata, ...) to the
+        # wrapped rewriter. ``__getattr__`` only fires for names not found normally,
+        # so the overrides above take precedence. Guard ``_rewriter`` itself to avoid
+        # infinite recursion before ``__init__`` has run.
+        if name == "_rewriter":
+            raise AttributeError(name)
+        return getattr(self._rewriter, name)
+
+
 class Mode:
     """A class that specifies the rewrites/optimizations used during function compilation.
 
@@ -377,9 +425,13 @@ def function_maker(self):
     @property
     def optimizer(self):
         if isinstance(self._optimizer, RewriteDatabaseQuery):
-            return self.optdb.query(self._optimizer)
+            rewriter = self.optdb.query(self._optimizer)
         else:
-            return self._optimizer
+            rewriter = self._optimizer
+        # Stamp this mode as the fgraph's active compile mode while rewriting, so
+        # inner-graph rewrites can recover the right linker even when the optimizer
+        # is run directly (outside ``FunctionMaker``, e.g. ``mode.optimizer.rewrite(fgraph)``).
+        return _ActiveModeGraphRewriter(rewriter, self)
 
     def get_linker_optimizer(self, linker, optimizer):
         if isinstance(linker, str) or linker is None:

pytensor/scalar/basic.py

@@ -4273,24 +4273,37 @@ def output_types(self, input_types):
         return self.outputs_type
 
     def make_node(self, *inputs):
-        if self.inputs_type == tuple(i.type for i in inputs):
+        inputs = [
+            inp if isinstance(inp, ScalarVariable) else as_scalar(inp) for inp in inputs
+        ]
+
+        if self.inputs_type == tuple(inp.type for inp in inputs):
             return super().make_node(*inputs)
+
+        if len(inputs) != self.nin:
+            raise ValueError("Number of inputs does not match expected")
+
+        # First try to coerce each input to its inner-graph input type.
+        try:
+            filtered_inputs = [
+                inner_inp.type.filter_variable(inp)
+                for inp, inner_inp in zip(inputs, self.inputs)
+            ]
+        except TypeError:
+            pass
         else:
-            # Make a new op whose inner graph is rebuilt on fresh inputs of the
-            # new types. The retype needs ``rebuild_strict=False`` (re-infers
-            # each node's output types), which in-place ``FunctionGraph``
-            # replacements cannot do, so thaw first and rebuild the mutable copy.
-            assert len(inputs) == self.nin
-            unfrozen_fgraph = self.fgraph.unfreeze()
-            new_inner_inputs = [i.type() for i in inputs]
-            new_outputs = clone_replace(
-                unfrozen_fgraph.outputs,
-                replace=dict(
-                    zip(unfrozen_fgraph.inputs, new_inner_inputs, strict=True)
-                ),
-                rebuild_strict=False,
-            )
-            return Composite(new_inner_inputs, new_outputs).make_node(*inputs)
+            return super().make_node(*filtered_inputs)
+
+        # The new input types are incompatible with the inner graph.
+        # Try to make a new inner graph with rebuild_strict=False.
+        unfrozen_fgraph = self.fgraph.unfreeze()
+        new_inner_inputs = [i.type() for i in inputs]
+        new_outputs = clone_replace(
+            unfrozen_fgraph.outputs,
+            replace=dict(zip(unfrozen_fgraph.inputs, new_inner_inputs)),
+            rebuild_strict=False,
+        )
+        return Composite(new_inner_inputs, new_outputs).make_node(*inputs)
 
     def perform(self, node, inputs, output_storage):
         outputs = self.py_perform_fn(*inputs)

pytensor/tensor/optimize.py

@@ -267,10 +267,10 @@ def prepare_node(
 
     def make_node(self, *inputs):
         assert len(inputs) == len(self.inner_inputs)
-        for input, inner_input in zip(inputs, self.inner_inputs):
-            assert input.type == inner_input.type, (
-                f"Input {input} does not match expected type {inner_input.type}"
-            )
+        inputs = [
+            inner_input.type.filter_variable(input)
+            for input, inner_input in zip(inputs, self.inner_inputs)
+        ]
 
         return Apply(self, inputs, [self.inner_inputs[0].type(), ps.bool("success")])
 

tests/compile/test_mode.py

@@ -8,6 +8,8 @@
 )
 from pytensor.configdefaults import config
 from pytensor.graph.features import NoOutputFromInplace
+from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.rewriting.basic import get_active_mode, graph_rewriter
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery, SequenceDB
 from pytensor.link.jax import JAXLinker
 from pytensor.tensor.math import dot, tanh
@@ -131,3 +133,24 @@ def test_predefined_modes_respected():
 
     default_mode_again = get_default_mode()
     assert not isinstance(default_mode_again.linker, JAXLinker)
+
+
+def test_optimizer_sets_active_compile_mode():
+    """``mode.optimizer.rewrite(fgraph)`` must expose ``mode`` as the active compile
+    mode (via ``get_active_mode``) while it runs, even outside ``FunctionMaker``, so
+    inner-graph rewrites recover the right backend. The mode is stamped only when
+    unset and restored afterwards.
+    """
+    seen = []
+
+    @graph_rewriter
+    def record_active_mode(fgraph):
+        seen.append(get_active_mode(fgraph))
+
+    mode = Mode(linker="py", optimizer=record_active_mode)
+    x = vector("x")
+    fg = FunctionGraph([x], [x])
+
+    mode.optimizer.rewrite(fg)
+    assert seen == [mode]
+    assert getattr(fg, "_compile_mode", None) is None

tests/tensor/test_optimize.py

@@ -103,6 +103,33 @@ def f(x, a, c):
     utt.verify_grad(f, [0.0, a_val, c_val], eps=1e-6)
 
 
+def test_minimize_replace_input_with_more_specific_type():
+    """Substituting an outer input of the minimize node with a compatible but
+    more specific type (a static shape ``(3,)`` in place of the symbolic ``(?,)``
+    of the minimized variable) must be accepted rather than rejected by a strict
+    ``input.type == inner_input.type`` check in ``make_node``. This mirrors the
+    Laplace/INLA init-point substitution ``graph_replace(x_star, {x: x0_init})``.
+    """
+    x = pt.vector("x")  # shape (?,)
+    a = pt.vector("a")  # shape (?,)
+    objective = ((x - a) ** 2).sum()
+    x_star, _ = minimize(objective, x)
+
+    x0_init = pt.ones(3)  # static shape (3,), more specific than x's (?,)
+    assert x0_init.type != x.type
+    # Before the fix this raised in ``ScipyWrapperOp.make_node``.
+    x_star_init = graph_replace(x_star, {x: x0_init})
+
+    f = function([a], x_star_init)
+    a_val = np.array([1.0, 2.0, 3.0], dtype=x.type.dtype)
+    np.testing.assert_allclose(
+        f(a_val),
+        a_val,
+        atol=1e-6,
+        rtol=1e-6,
+    )
+
+
 @pytest.mark.parametrize(
     "method, jac, hess",
     [