From 445d98808a9064e6f8f0d2a072e80866e4fb26f4 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Wed, 8 Jul 2026 16:27:03 +0200 Subject: [PATCH 1/2] Allow more specialized input types in MinimizeOp In Composite, only try to rebuild graph if types are not "filter"-compatible --- pytensor/scalar/basic.py | 45 ++++++++++++++++++++++------------- pytensor/tensor/optimize.py | 8 +++---- tests/tensor/test_optimize.py | 27 +++++++++++++++++++++ 3 files changed, 60 insertions(+), 20 deletions(-) diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py index ae03530fa6..759803932b 100644 --- a/pytensor/scalar/basic.py +++ b/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) diff --git a/pytensor/tensor/optimize.py b/pytensor/tensor/optimize.py index 532024d35b..eadfd66e3e 100644 --- a/pytensor/tensor/optimize.py +++ b/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")]) diff --git a/tests/tensor/test_optimize.py b/tests/tensor/test_optimize.py index 9c1d15fb65..cc673b1571 100644 --- a/tests/tensor/test_optimize.py +++ b/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", [ From edfb52e43042ea757f529f6c44c31572e2ba6258 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Wed, 8 Jul 2026 17:15:16 +0200 Subject: [PATCH 2/2] Move _compile_mode responsibility to mode.optimizer --- pytensor/compile/maker.py | 9 +----- pytensor/compile/mode.py | 56 ++++++++++++++++++++++++++++++++++++-- tests/compile/test_mode.py | 23 ++++++++++++++++ 3 files changed, 78 insertions(+), 10 deletions(-) diff --git a/pytensor/compile/maker.py b/pytensor/compile/maker.py index 7e9b103717..0073294012 100644 --- a/pytensor/compile/maker.py +++ b/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 diff --git a/pytensor/compile/mode.py b/pytensor/compile/mode.py index ecbf4be643..8e999511a2 100644 --- a/pytensor/compile/mode.py +++ b/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: diff --git a/tests/compile/test_mode.py b/tests/compile/test_mode.py index a4a699fe21..b43e25a272 100644 --- a/tests/compile/test_mode.py +++ b/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