pytensor 3.1.1__tar.gz → 3.1.2__tar.gz

{pytensor-3.1.1/pytensor.egg-info → pytensor-3.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.1.1
+Version: 3.1.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-3.1.1 → pytensor-3.1.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-07-03T07:15:26+0200",
+ "date": "2026-07-06T16:55:07+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "6b70d2cf2126abb036bbbe0f355e21957936eedf",
- "version": "3.1.1"
+ "full-revisionid": "47d3d0e3721d0488734b84051da78031081c7dc2",
+ "version": "3.1.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.1.1 → pytensor-3.1.2}/pytensor/compile/builders.py

@@ -8,8 +8,7 @@ from collections.abc import Callable, Sequence
 from copy import copy
 from functools import partial
 
-from pytensor.compile.io import In, Out
-from pytensor.compile.mode import get_mode
+from pytensor.compile.inner_function import HasInnerFunction
 from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.gradient import DisconnectedType, disconnected_type, grad, pushforward
 from pytensor.graph.basic import (
@@ -20,7 +19,7 @@ from pytensor.graph.basic import (
 )
 from pytensor.graph.fg import FrozenFunctionGraph, FunctionGraph
 from pytensor.graph.null_type import NullType
-from pytensor.graph.op import HasInnerGraph, Op, io_connection_pattern
+from pytensor.graph.op import Op, io_connection_pattern
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.traversal import graph_inputs
 from pytensor.graph.utils import MissingInputError
@@ -114,7 +113,7 @@ def construct_nominal_fgraph(
     return fgraph
 
 
-class OpFromGraph(Op, HasInnerGraph):
+class OpFromGraph(HasInnerFunction, Op):
     r"""Create an Op from inputs and outputs lists of variables.
 
     The signature is similar to :func:`pytensor.function` and the resulting Op's perform will do
@@ -840,50 +839,6 @@ class OpFromGraph(Op, HasInnerGraph):
 
         return ret
 
-    @property
-    def fn(self):
-        """Lazily compile the inner function graph."""
-        if getattr(self, "_fn", None) is not None:
-            return self._fn
-
-        # ``op.fgraph`` is already backend-optimized (inplace included): the
-        # ``ofg_inner_graph`` rewrite ran the backend optimizer on it during the
-        # outer compile. So we only need to link it. The linker forces
-        # ``minimum_compile`` back in via its ``required_rewrites``, and (for an
-        # inner graph) ``minimum_compile`` *is* that inner-graph rewrite -- so we
-        # exclude ``compile_inner_graph`` to stop it re-baking an already-baked
-        # graph. ``prepare_fgraph`` still inserts the boundary deepcopies; passing
-        # ``fgraph=`` avoids a re-clone. Unused inputs (e.g. rng, size) and
-        # internal-only inplace ops are expected and tolerated.
-        mode = (
-            get_mode(None)
-            .clone(optimizer="minimum_compile")
-            .excluding("compile_inner_graph")
-        )
-        unfrozen_fgraph = self.fgraph.unfreeze()
-        self._fn = mode.function_maker(
-            [In(inp, borrow=True) for inp in unfrozen_fgraph.inputs],
-            [Out(out, borrow=True) for out in unfrozen_fgraph.outputs],
-            mode,
-            fgraph=unfrozen_fgraph,
-            accept_inplace=True,
-            on_unused_input="ignore",
-        ).create()
-        self._fn.trust_input = True
-
-        return self._fn
-
-    @property
-    def inner_inputs(self):
-        # A list (not the frozen tuple) so callers that concatenate inner
-        # inputs/outputs keep list semantics. Read-only views of the immutable
-        # graph; manipulating them requires a fresh/unfrozen graph.
-        return list(self.fgraph.inputs)
-
-    @property
-    def inner_outputs(self):
-        return list(self.fgraph.outputs)
-
     def clone(self):
         # The inner graph is immutable (a frozen ``FunctionGraph``), so there is
         # nothing to deep-clone -- mirror ``Composite.clone``.

pytensor-3.1.2/pytensor/compile/inner_function.py

@@ -0,0 +1,84 @@
+"""Shared machinery for `Op`\\s that execute a lazily linked inner function."""
+
+from typing import TYPE_CHECKING
+
+from pytensor.compile.io import In, Out
+from pytensor.compile.mode import Mode
+from pytensor.graph.fg import FrozenFunctionGraph, FunctionGraph
+from pytensor.graph.op import HasInnerGraph
+from pytensor.graph.rewriting.basic import SequentialGraphRewriter
+from pytensor.link.basic import Linker
+
+
+if TYPE_CHECKING:
+    from pytensor.compile.function.types import Function
+
+
+def link_only_mode(linker: str | Linker) -> Mode:
+    """A `Mode` that links a graph without rewriting it at all.
+
+    The bare rewriter also bypasses the ``minimum_compile`` pass the linker
+    would otherwise force onto a database query.
+    """
+    return Mode(linker, SequentialGraphRewriter())
+
+
+class HasInnerFunction(HasInnerGraph):
+    """`HasInnerGraph` op whose ``perform`` runs a lazily linked inner function.
+
+    The frozen inner graph was already baked for the backend by the
+    ``compile_inner_graph`` rewrites during the outer compile, so linking it
+    needs no further rewrites (see `link_only_mode`).
+
+    The linker never comes from the config default mode: ``perform`` only runs
+    under the py/c backend family -- JIT backends (numba/jax/...) funcify
+    ``op.fgraph`` directly and never call ``perform`` -- so a JIT default must
+    not win, and the JIT inner-graph rewrites were never applied to this graph.
+    """
+
+    _fn = None
+    fgraph: FrozenFunctionGraph
+
+    def link_mode(self, impl: str | None) -> Mode:
+        """The `Mode` to link the inner function with, given a thunk ``impl``."""
+        return link_only_mode("cvm" if impl == "c" else "vm")
+
+    def link_fgraph(self, fgraph: FunctionGraph, mode: Mode) -> "Function":
+        """Link an already-baked inner ``fgraph`` under ``mode``, no rewrites."""
+        fn = mode.function_maker(
+            [In(inp) for inp in fgraph.inputs],
+            [Out(out) for out in fgraph.outputs],
+            mode,
+            fgraph=fgraph,
+            accept_inplace=True,
+            on_unused_input="ignore",
+        ).create()
+        fn.trust_input = True
+        return fn
+
+    def compile_fn(self, mode: Mode) -> "Function":
+        """Build the inner function under ``mode`` (override to massage the graph)."""
+        return self.link_fgraph(self.fgraph.unfreeze(), mode)
+
+    @property
+    def fn(self) -> "Function":
+        if self._fn is None:
+            self._fn = self.compile_fn(self.link_mode(None))
+        return self._fn
+
+    def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None):
+        if self._fn is None:
+            self._fn = self.compile_fn(self.link_mode(impl))
+        return super().make_thunk(
+            node, storage_map, compute_map, no_recycling, impl=impl
+        )
+
+    @property
+    def inner_inputs(self):
+        # Read-only views of the immutable inner graph, as lists so callers
+        # that concatenate inputs/outputs keep list semantics.
+        return list(self.fgraph.inputs)
+
+    @property
+    def inner_outputs(self):
+        return list(self.fgraph.outputs)

{pytensor-3.1.1 → pytensor-3.1.2}/pytensor/compile/mode.py

@@ -305,7 +305,7 @@ class Mode:
     def __init__(
         self,
         linker: str | Linker | None = None,
-        optimizer: str | RewriteDatabaseQuery = "default",
+        optimizer: str | RewriteDatabaseQuery | GraphRewriter = "default",
         db: RewriteDatabase = None,
     ):
         if linker is None:

{pytensor-3.1.1 → pytensor-3.1.2}/pytensor/link/numba/dispatch/elemwise.py

@@ -122,8 +122,12 @@ def scalar_in_place_fn_IntDiv(op, idx, res, arr):
 
 @scalar_in_place_fn.register(Maximum)
 def scalar_in_place_fn_Maximum(op, idx, res, arr):
+    # `arr != arr` catches NaN, which the comparison alone would drop; once the
+    # accumulator is NaN neither clause fires again, so NaN sticks (numpy
+    # semantics, matching the C backend). For integer dtypes LLVM folds the
+    # always-false clause away.
     return [
-        f"if {res}[{idx}] < {arr}:",
+        f"if {res}[{idx}] < {arr} or {arr} != {arr}:",
         CODE_TOKEN.INDENT,
         f"{res}[{idx}] = {arr}",
         CODE_TOKEN.DEDENT,
@@ -132,8 +136,9 @@ def scalar_in_place_fn_Maximum(op, idx, res, arr):
 
 @scalar_in_place_fn.register(Minimum)
 def scalar_in_place_fn_Minimum(op, idx, res, arr):
+    # See NaN comment in scalar_in_place_fn_Maximum
     return [
-        f"if {res}[{idx}] > {arr}:",
+        f"if {res}[{idx}] > {arr} or {arr} != {arr}:",
         CODE_TOKEN.INDENT,
         f"{res}[{idx}] = {arr}",
         CODE_TOKEN.DEDENT,
@@ -369,8 +374,11 @@ def create_multiaxis_reducer(
         else:
             identity = np.iinfo(acc_dtype).min
 
-    # Make sure it has the correct dtype
-    identity = getattr(np, acc_dtype.name)(identity)
+    # Make sure it has the correct dtype. Cast through ``astype`` rather than the
+    # scalar constructor so a negative identity (``-1`` for bitwise AND) wraps
+    # into an unsigned acc_dtype like C does instead of raising: ``np.uint64(-1)``
+    # is out of range, but ``-1`` cast to uint64 is all-ones -- the AND identity.
+    identity = np.asarray(identity).astype(acc_dtype)[()]
 
     kept_axes = [i for i in range(ndim) if i not in axes]
     n_kept = len(kept_axes)
@@ -840,7 +848,7 @@ def numba_funcify_CAReduce(op, node, **kwargs):
     )
     careduce_fn = numba_basic.numba_njit(careduce_py_fn, boundscheck=False)
 
-    cache_version = 3
+    cache_version = 4
     careduce_key = sha256(
         str(
             (

{pytensor-3.1.1 → pytensor-3.1.2}/pytensor/scan/op.py

@@ -59,6 +59,7 @@ from pytensor import tensor as pt
 from pytensor.compile.aliasing import add_supervisor_to_fgraph
 from pytensor.compile.builders import construct_nominal_fgraph, infer_shape
 from pytensor.compile.debug.profiling import register_profiler_printer
+from pytensor.compile.inner_function import HasInnerFunction, link_only_mode
 from pytensor.compile.io import In, Out
 from pytensor.compile.mode import Mode, get_mode
 from pytensor.configdefaults import config
@@ -76,7 +77,7 @@ from pytensor.graph.basic import (
 )
 from pytensor.graph.features import NoOutputFromInplace
 from pytensor.graph.fg import FrozenFunctionGraph, FunctionGraph
-from pytensor.graph.op import HasInnerGraph, Op, io_connection_pattern
+from pytensor.graph.op import Op, io_connection_pattern
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.traversal import graph_inputs
 from pytensor.graph.type import HasShape
@@ -835,7 +836,7 @@ class ScanMethodsMixin:
                     )
 
 
-class Scan(Op, ScanMethodsMixin, HasInnerGraph):
+class Scan(HasInnerFunction, Op, ScanMethodsMixin):
     r"""An `Op` implementing `for` and `while` loops.
 
     This `Op` has an "inner-graph" that represents the steps performed during
@@ -1542,12 +1543,26 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
         return wrapped_inputs, wrapped_outputs
 
-    @property
-    def fn(self):
-        """Lazily compile the inner function graph."""
-        if getattr(self, "_fn", None) is not None:
-            return self._fn
+    def link_mode(self, impl):
+        # ``self.mode`` is a deprecated per-op override, respected when given:
+        # its linker is used as-is. Otherwise the default py/c rule applies,
+        # except FAST_COMPILE forces the pure-python VM.
+        mode = self.mode
+        if mode in (None, "FAST_RUN"):
+            return super().link_mode(impl)
+        if mode == "FAST_COMPILE":
+            return link_only_mode(VMLinker(use_cloop=False, c_thunks=False))
+        linker = get_mode(mode).clone(link_kwargs=dict(allow_gc=self.allow_gc)).linker
+        # Scan's python/cython perform sets preallocated MIT-MOT updates through
+        # the VM, which only a VMLinker provides.
+        if any(self.mitmots_preallocated) and not isinstance(linker, VMLinker):
+            raise NotImplementedError(
+                "Python/Cython implementation of Scan with preallocated MIT-MOT "
+                f"outputs requires a VMLinker, got {linker}"
+            )
+        return link_only_mode(linker)
 
+    def compile_fn(self, mode):
         # Compile a throwaway copy of the (already math-optimized) inner graph.
         # The canonical inner graph is immutable; linking setup (MIT-MOT update
         # wrapping, supervisor) and any inplace happen on this transient.
@@ -1565,45 +1580,10 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         elif self.profile:
             profile = self.profile
 
-        # ``inner_fgraph`` is already backend-optimized (inplace included):
-        # ``scan_inner_graph`` ran the backend optimizer on it during the outer
-        # compile. So we only need to link it (``prepare_fgraph`` still inserts the
-        # boundary deepcopies). The linker forces ``minimum_compile`` back in via
-        # its ``required_rewrites``, and (for an inner graph) ``minimum_compile``
-        # *is* that inner-graph rewrite -- so we exclude ``compile_inner_graph`` to
-        # stop it re-baking an already-baked graph. Only the linker choice depends
-        # on ``self.mode``.
-        mode = self.mode
-        if mode in (None, "FAST_RUN"):
-            mode_instance = Mode("cvm", "minimum_compile").excluding(
-                "compile_inner_graph"
-            )
-        elif mode == "FAST_COMPILE":
-            mode_instance = Mode(
-                VMLinker(use_cloop=False, c_thunks=False), "minimum_compile"
-            ).excluding("compile_inner_graph")
-        else:
-            mode_instance = (
-                get_mode(mode)
-                .clone(
-                    optimizer="minimum_compile",
-                    link_kwargs=dict(allow_gc=self.allow_gc),
-                    message=f"{self.name or 'Scan'} sub profile",
-                )
-                .excluding("compile_inner_graph")
-            )
-            # Scan python and cython perform relies on the VM being able to set updates for preallocated MIT-MOT,
-            # which only the VMs produced by VMLinker do
-            if any(self.mitmots_preallocated) and not isinstance(
-                mode_instance.linker, VMLinker
-            ):
-                raise NotImplementedError(
-                    f"Python/Cython implementation of Scan with preallocated MIT-MOT outputs requires a VMLinker, got {mode_instance.linker}"
-                )
-        self._fn = mode_instance.function_maker(
+        return mode.function_maker(
             wrapped_inputs,
             wrapped_outputs,
-            mode=mode_instance,
+            mode=mode,
             # The (already-optimized) inner graph may carry inplace ops baked in
             # by scan_inner_graph; prepare_fgraph has already attached the
             # DestroyHandler + Supervisor, so accept them here.
@@ -1613,19 +1593,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             fgraph=inner_fgraph,
         ).create()
 
-        return self._fn
-
-    @property
-    def inner_inputs(self):
-        # A list (not the frozen tuple) so the many ``inner_*`` slicing helpers
-        # and their callers keep list semantics. These are read-only views of the
-        # immutable graph; rewrites that rebuild a Scan must ``unfreeze`` first.
-        return list(self.fgraph.inputs)
-
-    @property
-    def inner_outputs(self):
-        return list(self.fgraph.outputs)
-
     def clone(self) -> "Scan":
         # The inner graph is immutable (a frozen ``FunctionGraph``), so there is
         # nothing to deep-clone -- mirror ``Composite.clone``.
@@ -1686,6 +1653,11 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         # coherent
         self.validate_inner_graph()
 
+        # Lazily link the inner function for this thunk's backend (see
+        # ``HasInnerFunction``); Scan drives the resulting VM itself below.
+        if self._fn is None:
+            self._fn = self.compile_fn(self.link_mode(impl))
+
         # Setting up all my variables in what I believe is a more Cython
         # friendly form
 

{pytensor-3.1.1 → pytensor-3.1.2}/pytensor/tensor/optimize.py

@@ -6,15 +6,13 @@ from functools import singledispatch
 import numpy as np
 
 import pytensor.scalar as ps
-from pytensor.compile.maker import function
-from pytensor.compile.mode import get_mode
+from pytensor.compile.inner_function import HasInnerFunction
 from pytensor.gradient import DisconnectedType, grad, jacobian
 from pytensor.graph.basic import Apply, Constant
 from pytensor.graph.fg import FrozenFunctionGraph, FunctionGraph
 from pytensor.graph.null_type import NullType
 from pytensor.graph.op import (
     ComputeMapType,
-    HasInnerGraph,
     Op,
     StorageMapType,
     io_connection_pattern,
@@ -164,7 +162,37 @@ def _find_optimization_parameters(
     ]
 
 
-class ScipyWrapperOp(Op, HasInnerGraph):
+def _tensorize_scalar_parameters(
+    objective: TensorVariable, args: list[Variable]
+) -> tuple[TensorVariable, list[Variable], list[Variable]]:
+    """Replace ``ScalarType`` parameters with equivalent ``TensorType`` inputs.
+
+    scipy hands every value to the objective as a numpy array, and stores/forwards
+    a ``ScalarType`` value as a bare python float -- which no C thunk accepts under
+    ``trust_input``. Re-expressing such a parameter as a 0-d ``TensorType`` input
+    (reconstructed inside the graph with ``scalar_from_tensor``) makes the op
+    boundary array-typed, so scipy passes a preserved 0-d ndarray. Returns the
+    rewritten objective, the inner-graph parameters, and the matching outer inputs.
+    """
+    replacements: dict[Variable, Variable] = {}
+    inner_args: list[Variable] = []
+    outer_args: list[Variable] = []
+    for arg in args:
+        if isinstance(arg.type, ScalarType):
+            inner = tensor(name=arg.name, shape=(), dtype=arg.type.dtype)
+            replacements[arg] = scalar_from_tensor(inner)
+            inner_args.append(inner)
+            outer_args.append(tensor_from_scalar(arg))
+        else:
+            inner_args.append(arg)
+            outer_args.append(arg)
+
+    if replacements:
+        [objective] = graph_replace([objective], replacements)
+    return objective, inner_args, outer_args
+
+
+class ScipyWrapperOp(HasInnerFunction, Op):
     """Shared logic for scipy optimization ops.
 
     The inner graph is held frozen (immutable) as ``self.fgraph``, so the
@@ -178,57 +206,12 @@ class ScipyWrapperOp(Op, HasInnerGraph):
     # of the same type for eq/hash. Subclasses override.
     _scipy_props: tuple[str, ...] = ()
 
-    def build_fn(self):
-        """
-        This is overloaded because scipy converts scalar inputs to lists, changing the return type. The
-        wrapper function logic is there to handle this.
-        """
-        fgraph = self.fgraph.unfreeze()
-        # ``optimize_inner_graph`` already baked this graph for the active backend
-        # (inplace included), exactly like ``OpFromGraph``. So we only link it --
-        # ``minimum_compile`` excluding ``compile_inner_graph`` (the rewrite that
-        # already ran); ``prepare_fgraph`` still inserts the boundary deepcopies.
-        # ``accept_inplace`` admits the baked inplace ops; backend optimization can
-        # leave a declared input unused (e.g. a folded core-shape vector), hence
-        # ``on_unused_input="ignore"``.
-        self._fn = fn = function(
-            fgraph.inputs,
-            fgraph.outputs,
-            mode=get_mode(None)
-            .clone(optimizer="minimum_compile")
-            .excluding("compile_inner_graph"),
-            accept_inplace=True,
-            trust_input=True,
-            on_unused_input="ignore",
-        )
-
-        # Do this reassignment to see the compiled graph in the dprint
-        # self.fgraph = fn.maker.fgraph
-
-        self._fn_wrapped = LRUCache1(fn)
-
-    @property
-    def fn(self):
-        if self._fn is None:
-            self.build_fn()
-        return self._fn
-
     @property
     def fn_wrapped(self):
         if self._fn_wrapped is None:
-            self.build_fn()
+            self._fn_wrapped = LRUCache1(self.fn)
         return self._fn_wrapped
 
-    @property
-    def inner_inputs(self):
-        # A list (not the frozen tuple) so callers that concatenate inner
-        # inputs/outputs keep list semantics.
-        return list(self.fgraph.inputs)
-
-    @property
-    def inner_outputs(self):
-        return list(self.fgraph.outputs)
-
     def _prop_values(self):
         return tuple(getattr(self, name) for name in self._scipy_props)
 
@@ -280,7 +263,7 @@ class ScipyWrapperOp(Op, HasInnerGraph):
         impl: str | None,
     ):
         """Trigger the compilation of the inner fgraph so it shows in the dprint before the first call"""
-        self.build_fn()
+        self.fn
 
     def make_node(self, *inputs):
         assert len(inputs) == len(self.inner_inputs)
@@ -305,7 +288,7 @@ def rewrite_optimize_inner_graph(linker, op, node, inner, *, mode):
 
 
 class ScipyScalarWrapperOp(ScipyWrapperOp):
-    def build_fn(self):
+    def compile_fn(self, mode):
         # We need to adjust the graph to work with what scipy will be passing into the inner function --
         # always scalar array of float64 type
         fgraph = self.fgraph.unfreeze()
@@ -315,24 +298,10 @@ class ScipyScalarWrapperOp(ScipyWrapperOp):
 
         new_outputs = graph_replace(fgraph.outputs, {x: new_x})
 
-        # See ``ScipyWrapperOp.build_fn``: the graph is already baked, so link
-        # it (the scipy boundary wrapping above links alongside the baked body).
-        self._fn = fn = function(
-            [new_root_x, *args],
-            new_outputs,
-            mode=get_mode(None)
-            .clone(optimizer="minimum_compile")
-            .excluding("compile_inner_graph"),
-            accept_inplace=True,
-            trust_input=True,
-            on_unused_input="ignore",
+        return self.link_fgraph(
+            FunctionGraph([new_root_x, *args], new_outputs, clone=False), mode
         )
 
-        # Do this reassignment to see the compiled graph in the dprint
-        # self.fgraph = fn.maker.fgraph
-
-        self._fn_wrapped = LRUCache1(fn)
-
     def compute_implicit_gradients(
         self,
         x_star: TensorVariable,
@@ -407,34 +376,21 @@ class ScipyScalarWrapperOp(ScipyWrapperOp):
 
 
 class ScipyVectorWrapperOp(ScipyWrapperOp):
-    def build_fn(self):
+    def compile_fn(self, mode):
         # We need to adjust the graph to work with what scipy will be passing into the inner function --
         # always a vector array with size of at least 1
         if self.inner_inputs[0].type.shape != ():
-            return super().build_fn()
+            return super().compile_fn(mode)
 
         fgraph = self.fgraph.unfreeze()
         x, *args = fgraph.inputs
         new_root_x = x[None].type()
         new_x = new_root_x.squeeze()
         new_outputs = graph_replace(fgraph.outputs, {x: new_x})
-        # See ``ScipyWrapperOp.build_fn``: the graph is already baked, so link
-        # it (the scipy boundary wrapping above links alongside the baked body).
-        self._fn = fn = function(
-            [new_root_x, *args],
-            new_outputs,
-            mode=get_mode(None)
-            .clone(optimizer="minimum_compile")
-            .excluding("compile_inner_graph"),
-            accept_inplace=True,
-            trust_input=True,
-            on_unused_input="ignore",
+        return self.link_fgraph(
+            FunctionGraph([new_root_x, *args], new_outputs, clone=False), mode
         )
 
-        # Do this reassignment to see the compiled graph in the dprint
-        # self.fgraph = fn.maker.fgraph
-        self._fn_wrapped = LRUCache1(fn)
-
     def compute_implicit_gradients(
         self,
         x_star: TensorVariable,
@@ -698,6 +654,7 @@ def minimize_scalar(
         value, based on the requested convergence criteria.
     """
     args = _find_optimization_parameters(objective, x)
+    objective, args, outer_args = _tensorize_scalar_parameters(objective, args)
 
     minimize_scalar_op = MinimizeScalarOp(
         x,
@@ -707,7 +664,7 @@ def minimize_scalar(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    solution, success = minimize_scalar_op(x, *args)
+    solution, success = minimize_scalar_op(x, *outer_args)
 
     return solution, success
 
@@ -905,6 +862,7 @@ def minimize(
 
     packed_input, packed_shapes, objective = pack_inputs_of_objective(objective, x)
     args = _find_optimization_parameters(objective, packed_input)
+    objective, args, outer_args = _tensorize_scalar_parameters(objective, args)
 
     minimize_op = MinimizeOp(
         packed_input,
@@ -917,7 +875,7 @@ def minimize(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    solution, success = minimize_op(packed_input, *args)
+    solution, success = minimize_op(packed_input, *outer_args)
 
     if packed_shapes is not None:
         solution = unpack(solution, packed_shapes)
@@ -1065,6 +1023,7 @@ def root_scalar(
         Boolean indicating whether the root-finding was successful. If True, the solution is a root of the equation
     """
     args = _find_optimization_parameters(equation, variable)
+    equation, args, outer_args = _tensorize_scalar_parameters(equation, args)
 
     root_scalar_op = RootScalarOp(
         variable,
@@ -1076,7 +1035,7 @@ def root_scalar(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    solution, success = root_scalar_op(variable, *args)
+    solution, success = root_scalar_op(variable, *outer_args)
 
     return solution, success
 
@@ -1139,7 +1098,7 @@ class RootOp(ScipyVectorWrapperOp):
         )
         return f"{self.__class__.__name__}({str_args})"
 
-    def build_fn(self):
+    def compile_fn(self, mode):
         fgraph = self.fgraph.unfreeze()
         variables, *args = fgraph.inputs
         outputs = fgraph.outputs
@@ -1157,24 +1116,10 @@ class RootOp(ScipyVectorWrapperOp):
 
             new_outputs = graph_replace(outputs, {variables: new_variables})
 
-        # See ``ScipyWrapperOp.build_fn``: the graph is already baked, so link
-        # it (the scipy boundary wrapping above links alongside the baked body).
-        self._fn = fn = function(
-            [new_root_variables, *args],
-            new_outputs,
-            mode=get_mode(None)
-            .clone(optimizer="minimum_compile")
-            .excluding("compile_inner_graph"),
-            accept_inplace=True,
-            trust_input=True,
-            on_unused_input="ignore",
+        return self.link_fgraph(
+            FunctionGraph([new_root_variables, *args], new_outputs, clone=False), mode
         )
 
-        # Do this reassignment to see the compiled graph in the dprint
-        # self.fgraph = fn.maker.fgraph
-
-        self._fn_wrapped = LRUCache1(fn)
-
     def perform(self, node, inputs, outputs):
         global optimize
         if optimize is None:
@@ -1260,6 +1205,7 @@ def root(
         equations, variables
     )
     args = _find_optimization_parameters(equations, packed_variables)
+    equations, args, outer_args = _tensorize_scalar_parameters(equations, args)
 
     root_op = RootOp(
         packed_variables,
@@ -1271,7 +1217,7 @@ def root(
         use_vectorized_jac=use_vectorized_jac,
     )
 
-    solution, success = root_op(packed_variables, *args)
+    solution, success = root_op(packed_variables, *outer_args)
     if packed_shapes is not None:
         solution = unpack(solution, packed_shapes)