pytensor 2.35.1__tar.gz → 2.36.0__tar.gz

{pytensor-2.35.1/pytensor.egg-info → pytensor-2.36.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.35.1
+Version: 2.36.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -24,12 +24,14 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Requires-Python: <3.14,>=3.11
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: <3.15,>=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 Requires-Dist: setuptools>=59.0.0
 Requires-Dist: scipy<2,>=1
 Requires-Dist: numpy>=2.0
+Requires-Dist: numba<1,>0.57
 Requires-Dist: filelock>=3.15
 Requires-Dist: etuples
 Requires-Dist: logical-unification

{pytensor-2.35.1 → pytensor-2.36.0}/doc/extending/creating_a_numba_jax_op.rst

@@ -228,7 +228,7 @@ Here's an example for :class:`DimShuffle`:
                     # E       No match.
                     # ...(on this line)...
                     # E           shuffle_shape = res.shape[: len(shuffle)]
-                    @numba_basic.numba_njit(inline="always")
+                    @numba_basic.numba_njit
                     def dimshuffle(x):
                         return dimshuffle_inner(np.asarray(x), shuffle)
 

{pytensor-2.35.1 → pytensor-2.36.0}/doc/gallery/autodiff/vector_jacobian_product.ipynb

@@ -389,7 +389,7 @@
     "1 & 1 & 0 \\\\\n",
     "1 & 1 & 1 \\\\\n",
     "\\end{pmatrix} = \\begin{pmatrix}\n",
-    "v_1 + v2 + v 3 \\\\ v_2 + v_3 \\\\ v_3\n",
+    "v_1 + v_2 + v_3 \\\\ v_2 + v_3 \\\\ v_3\n",
     "\\end{pmatrix}\n",
     "$$"
    ]
@@ -661,7 +661,7 @@
    "source": [
     "$$\n",
     "v^T \\cdot J = \\begin{pmatrix}\n",
-    "v_1 \\\\ v_2 \\\\ v_3 \\\\ v4\n",
+    "v_1 \\\\ v_2 \\\\ v_3 \\\\ v_4\n",
     "\\end{pmatrix}^T \\cdot \\begin{pmatrix}\n",
     "x_1 & 0 \\\\\n",
     "x_2 & x_1 \\\\\n",

{pytensor-2.35.1 → pytensor-2.36.0}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "pytensor"
 dynamic = ['version']
-requires-python = ">=3.11,<3.14"
+requires-python = ">=3.11,<3.15"
 authors = [{ name = "pymc-devs", email = "pymc.devs@gmail.com" }]
 description = "Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs."
 readme = "README.rst"
@@ -33,6 +33,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
 ]
 
 keywords = [
@@ -49,6 +50,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
     "numpy>=2.0",
+    "numba>0.57,<1",
     "filelock>=3.15",
     "etuples",
     "logical-unification",
@@ -120,6 +122,9 @@ tag_prefix = "rel-"
 addopts = "--durations=50 --doctest-modules --ignore=pytensor/link --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/ipython.py"
 testpaths = ["pytensor/", "tests/"]
 xfail_strict = true
+filterwarnings =[
+    'ignore:^Numba will use object mode to run.*perform method\.:UserWarning',
+]
 
 [tool.ruff]
 line-length = 88
@@ -150,7 +155,7 @@ lines-after-imports = 2
 "pytensor/misc/check_duplicate_key.py" = ["T201"]
 "pytensor/misc/check_blas.py" = ["T201"]
 "pytensor/bin/pytensor_cache.py" = ["T201"]
-# For the tests we skip because `pytest.importorskip` is used:
+# For the tests we skip `E402` because `pytest.importorskip` is used:
 "tests/link/jax/test_scalar.py" = ["E402"]
 "tests/link/jax/test_tensor_basic.py" = ["E402"]
 "tests/link/numba/test_basic.py" = ["E402"]
@@ -182,12 +187,12 @@ files = ["pytensor", "tests"]
 build = "*"
 # Uncomment to skip builds that compile but fail when trying to test (maybe due to incompatibility with runner)
 # archs = ["auto64"]
-# Disable any-platform (pp*), and 32-bit builds.
+# Disable any-platform (pp*), 32-bit builds, and free-threaded builds (cp*t-*)
 # Additional options to consider: "*musllinux*"
-skip = ["pp*", "*-win32", "*-manylinux_i686"]
+skip = ["pp*", "*-win32", "*-manylinux_i686", "cp*t-*"]
 build-frontend = "build"
 test-command = 'python -c "import pytensor; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version())"'
-test-skip = ["cp37*", "*musllinux*", "*win32*", "*i686*"]
+test-skip = ["*musllinux*", "*i686*"]
 
 # Testing seems to be running into issues locating libs where expected
 # test-requires = ["pytest", "numba", "jax", "jaxlib"]

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-10-20T17:04:33+0200",
+ "date": "2025-12-15T22:58:11+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "5547eb086c6f281621fb81f2ef93e041ee05aece",
- "version": "2.35.1"
+ "full-revisionid": "6319fac84ebc71ad079c5986c6c302a52e254198",
+ "version": "2.36.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/bin/pytensor_cache.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import shutil
 import sys
 from pathlib import Path
 
@@ -74,7 +75,10 @@ def main():
                     'You can also call "pytensor-cache purge" to '
                     "remove everything from that directory."
                 )
-                _logger.debug(f"Remaining elements ({len(items)}): {', '.join(items)}")
+                _logger.debug(f"Remaining elements ({len(items)}): {items}")
+            numba_cache_dir: Path = config.base_compiledir / "numba"
+            shutil.rmtree(numba_cache_dir, ignore_errors=True)
+
         elif sys.argv[1] == "list":
             pytensor.compile.compiledir.print_compiledir_content()
         elif sys.argv[1] == "cleanup":
@@ -86,6 +90,8 @@ def main():
             print("Lock successfully removed!")
         elif sys.argv[1] == "purge":
             pytensor.compile.compiledir.compiledir_purge()
+            numba_cache_dir: Path = config.base_compiledir / "numba"
+            shutil.rmtree(numba_cache_dir, ignore_errors=True)
         elif sys.argv[1] == "basecompiledir":
             # Simply print the base_compiledir
             print(pytensor.config.base_compiledir)

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/__init__.py

@@ -17,8 +17,8 @@ from pytensor.compile.function.types import (
 )
 from pytensor.compile.io import In, Out, SymbolicInput, SymbolicOutput
 from pytensor.compile.mode import (
+    CVM,
     FAST_COMPILE,
-    FAST_RUN,
     JAX,
     NUMBA,
     OPT_FAST_COMPILE,
@@ -33,6 +33,7 @@ from pytensor.compile.mode import (
     PYTORCH,
     AddDestroyHandler,
     AddFeatureOptimizer,
+    C,
     Mode,
     PrintCurrentFunctionGraph,
     get_default_mode,

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/builders.py

@@ -4,6 +4,7 @@ import warnings
 from collections.abc import Callable, Sequence
 from copy import copy
 from functools import partial
+from itertools import chain
 from typing import Union, cast
 
 from pytensor.compile.function import function
@@ -47,11 +48,15 @@ def infer_shape(outs, inputs, input_shapes):
             assert len(inp_shp) == inp.type.ndim
 
     shape_feature = ShapeFeature()
-    shape_feature.on_attach(FunctionGraph([], []))
+    fgraph = FunctionGraph([], [], features=[shape_feature])
+    for v in chain.from_iterable(s for s in input_shapes if s is not None):
+        # Import input_shape nodes, as for some graphs ShapeFeature assumes these were seen before
+        if (node := v.owner) is not None:
+            fgraph.import_node(node, import_missing=True)
 
     # Initialize shape_of with the input shapes
     for inp, inp_shp in zip(inputs, input_shapes, strict=True):
-        shape_feature.set_shape(inp, inp_shp)
+        shape_feature.set_shape(inp, inp_shp, override=True)
 
     def local_traverse(out):
         """

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/debugmode.py

@@ -1331,7 +1331,11 @@ default_make_thunk = [get_unbound_function(COp.make_thunk)]
 # the external requirements of the .linker attribute of a mode
 # 1) it's a class instance
 # 2) it a has a .clone() method
+# 3) it has required_rewrites and incompatible_rewrites class attributes
 class _DummyLinker:
+    required_rewrites = ()
+    incompatible_rewrites = ()
+
     # This is not a real linker anyway
     def clone(self, allow_gc=None):
         return self

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/function/pfunc.py

@@ -453,7 +453,6 @@ def pfunc(
     inputs, cloned_outputs = construct_pfunc_ins_and_outs(
         params,
         outputs,
-        mode,
         updates,
         givens,
         no_default_updates,
@@ -479,7 +478,6 @@ def pfunc(
 def construct_pfunc_ins_and_outs(
     params,
     outputs=None,
-    mode=None,
     updates=None,
     givens=None,
     no_default_updates=False,
@@ -546,26 +544,27 @@ def construct_pfunc_ins_and_outs(
                 "variables in the inputs list."
             )
 
-    # Check that we are not using `givens` to replace input variables, because
-    # this typically does nothing, contrary to what one may expect.
-    in_var_set = set(in_variables)
-    try:
-        givens_pairs = list(givens.items())
-    except AttributeError:
-        givens_pairs = givens
-    for x, y in givens_pairs:
-        if x in in_var_set:
-            raise RuntimeError(
-                f"You are trying to replace variable '{x}' through the "
-                "`givens` parameter, but this variable is an input to your "
-                "function. Replacing inputs is currently forbidden because it "
-                "has no effect. One way to modify an input `x` to a function "
-                "evaluating f(x) is to define a new input `y` and use "
-                "`pytensor.function([y], f(x), givens={x: g(y)})`. Another "
-                "solution consists in using `pytensor.clone_replace`, e.g. like this: "
-                "`pytensor.function([x], "
-                "pytensor.clone_replace(f(x), replace={x: g(x)}))`."
-            )
+    if givens:
+        # Check that we are not using `givens` to replace input variables, because
+        # this typically does nothing, contrary to what one may expect.
+        in_var_set = set(in_variables)
+        try:
+            givens_pairs = list(givens.items())
+        except AttributeError:
+            givens_pairs = givens
+        for x, y in givens_pairs:
+            if x in in_var_set:
+                raise RuntimeError(
+                    f"You are trying to replace variable '{x}' through the "
+                    "`givens` parameter, but this variable is an input to your "
+                    "function. Replacing inputs is currently forbidden because it "
+                    "has no effect. One way to modify an input `x` to a function "
+                    "evaluating f(x) is to define a new input `y` and use "
+                    "`pytensor.function([y], f(x), givens={x: g(y)})`. Another "
+                    "solution consists in using `pytensor.clone_replace`, e.g. like this: "
+                    "`pytensor.function([x], "
+                    "pytensor.clone_replace(f(x), replace={x: g(x)}))`."
+                )
 
     if not fgraph:
         # Extend the outputs with the updates on input variables so they are

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/io.py

@@ -95,7 +95,7 @@ class SymbolicInput:
         self.implicit = implicit
 
     def __str__(self):
-        if self.update:
+        if self.update is not None:
             return f"In({self.variable} -> {self.update})"
         else:
             return f"In({self.variable})"

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/mode.py

@@ -5,7 +5,7 @@ WRITEME
 
 import logging
 import warnings
-from typing import Literal
+from typing import Any, Literal
 
 from pytensor.compile.function.types import Supervisor
 from pytensor.configdefaults import config
@@ -62,23 +62,17 @@ def register_linker(name, linker):
     predefined_linkers[name] = linker
 
 
-# If a string is passed as the optimizer argument in the constructor
-# for Mode, it will be used as the key to retrieve the real optimizer
-# in this dictionary
-exclude = []
-if not config.cxx:
-    exclude = ["cxx_only"]
-OPT_NONE = RewriteDatabaseQuery(include=[], exclude=exclude)
+OPT_NONE = RewriteDatabaseQuery(include=[])
 # Minimum set of rewrites needed to evaluate a function. This is needed for graphs with "dummy" Operations
-OPT_MINIMUM = RewriteDatabaseQuery(include=["minimum_compile"], exclude=exclude)
+OPT_MINIMUM = RewriteDatabaseQuery(include=["minimum_compile"])
 # Even if multiple merge optimizer call will be there, this shouldn't
 # impact performance.
-OPT_MERGE = RewriteDatabaseQuery(include=["merge"], exclude=exclude)
-OPT_FAST_RUN = RewriteDatabaseQuery(include=["fast_run"], exclude=exclude)
+OPT_MERGE = RewriteDatabaseQuery(include=["merge"])
+OPT_FAST_RUN = RewriteDatabaseQuery(include=["fast_run"])
 OPT_FAST_RUN_STABLE = OPT_FAST_RUN.requiring("stable")
 
-OPT_FAST_COMPILE = RewriteDatabaseQuery(include=["fast_compile"], exclude=exclude)
-OPT_STABILIZE = RewriteDatabaseQuery(include=["fast_run"], exclude=exclude)
+OPT_FAST_COMPILE = RewriteDatabaseQuery(include=["fast_compile"])
+OPT_STABILIZE = RewriteDatabaseQuery(include=["fast_run"])
 OPT_STABILIZE.position_cutoff = 1.5000001
 OPT_NONE.name = "OPT_NONE"
 OPT_MINIMUM.name = "OPT_MINIMUM"
@@ -316,6 +310,8 @@ class Mode:
     ):
         if linker is None:
             linker = config.linker
+        if isinstance(linker, str) and linker == "auto":
+            linker = "cvm" if config.cxx else "vm"
         if isinstance(optimizer, str) and optimizer == "default":
             optimizer = config.optimizer
 
@@ -352,7 +348,14 @@ class Mode:
         if isinstance(optimizer, str) or optimizer is None:
             optimizer = predefined_optimizers[optimizer]
         if isinstance(optimizer, RewriteDatabaseQuery):
+            # TODO: From the __init__ signature this should always be the case
+            # But some tests and internal logic allow passing a GraphRewriter directly as optimizer
+            # Cleanup!
             self.provided_optimizer = optimizer
+            if r := linker.required_rewrites:
+                optimizer = optimizer.including(*r)
+            if r := linker.incompatible_rewrites:
+                optimizer = optimizer.excluding(*r)
         self._optimizer = optimizer
         self.call_time = 0
         self.fn_time = 0
@@ -365,14 +368,13 @@ class Mode:
             f"optdb={self.optdb})"
         )
 
-    def __get_optimizer(self):
+    @property
+    def optimizer(self):
         if isinstance(self._optimizer, RewriteDatabaseQuery):
             return self.optdb.query(self._optimizer)
         else:
             return self._optimizer
 
-    optimizer = property(__get_optimizer)
-
     def get_linker_optimizer(self, linker, optimizer):
         if isinstance(linker, str) or linker is None:
             linker = predefined_linkers[linker]
@@ -445,92 +447,51 @@ class Mode:
         return new_mode
 
 
-# If a string is passed as the mode argument in function or
-# FunctionMaker, the Mode will be taken from this dictionary using the
-# string as the key
-# Use VM_linker to allow lazy evaluation by default.
-FAST_COMPILE = Mode(
-    VMLinker(use_cloop=False, c_thunks=False),
-    RewriteDatabaseQuery(include=["fast_compile", "py_only"]),
-)
-if config.cxx:
-    FAST_RUN = Mode("cvm", "fast_run")
-else:
-    FAST_RUN = Mode(
-        "vm",
-        RewriteDatabaseQuery(include=["fast_run", "py_only"]),
-    )
+C = Mode("c", "fast_run")
+CVM = Mode("cvm", "fast_run")
+VM = (Mode("vm", "fast_run"),)
 
 NUMBA = Mode(
     NumbaLinker(),
-    RewriteDatabaseQuery(
-        include=["fast_run", "numba"],
-        exclude=[
-            "cxx_only",
-            "BlasOpt",
-            "local_careduce_fusion",
-            "scan_save_mem_prealloc",
-        ],
-    ),
+    RewriteDatabaseQuery(include=["fast_run", "numba"]),
 )
 
 JAX = Mode(
     JAXLinker(),
-    RewriteDatabaseQuery(
-        include=["fast_run", "jax"],
-        exclude=[
-            "cxx_only",
-            "BlasOpt",
-            "fusion",
-            "inplace",
-            "scan_save_mem_prealloc",
-            # There are specific variants for the LU decompositions supported by JAX
-            "reuse_lu_decomposition_multiple_solves",
-            "scan_split_non_sequence_lu_decomposition_solve",
-        ],
-    ),
+    RewriteDatabaseQuery(include=["fast_run", "jax"]),
 )
 PYTORCH = Mode(
     PytorchLinker(),
-    RewriteDatabaseQuery(
-        include=["fast_run"],
-        exclude=[
-            "cxx_only",
-            "BlasOpt",
-            "fusion",
-            "inplace",
-            "scan_save_mem_prealloc",
-            "reuse_lu_decomposition_multiple_solves",
-            "scan_split_non_sequence_lu_decomposition_solve",
-        ],
-    ),
+    RewriteDatabaseQuery(include=["fast_run"]),
 )
 
 MLX = Mode(
     MLXLinker(),
-    RewriteDatabaseQuery(
-        include=["fast_run"],
-        exclude=[
-            "cxx_only",
-            "BlasOpt",
-            "fusion",
-            "inplace",
-            "scan_save_mem_prealloc",
-        ],
-    ),
+    RewriteDatabaseQuery(include=["fast_run"]),
+)
+
+FAST_COMPILE = Mode(
+    VMLinker(use_cloop=False, c_thunks=False),
+    RewriteDatabaseQuery(include=["fast_compile", "py_only"]),
 )
 
+fast_run_linkers_to_mode = {
+    "cvm": CVM,
+    "vm": VM,
+    "numba": NUMBA,
+}
 
 predefined_modes = {
     "FAST_COMPILE": FAST_COMPILE,
-    "FAST_RUN": FAST_RUN,
+    "C": C,
+    "CVM": CVM,
     "JAX": JAX,
     "NUMBA": NUMBA,
     "PYTORCH": PYTORCH,
     "MLX": MLX,
 }
 
-_CACHED_RUNTIME_MODES: dict[str, Mode] = {}
+_CACHED_RUNTIME_MODES: dict[Any, Mode] = {}
 
 
 def get_mode(orig_string):
@@ -548,10 +509,20 @@ def get_mode(orig_string):
     if upper_string in predefined_modes:
         return predefined_modes[upper_string]
 
+    if upper_string == "FAST_RUN":
+        linker = config.linker
+        if linker == "auto":
+            return CVM if config.cxx else VM
+        return fast_run_linkers_to_mode[linker]
+
     global _CACHED_RUNTIME_MODES
 
-    if upper_string in _CACHED_RUNTIME_MODES:
-        return _CACHED_RUNTIME_MODES[upper_string]
+    cache_key = ("MODE", config.linker) if upper_string == "MODE" else upper_string
+
+    try:
+        return _CACHED_RUNTIME_MODES[cache_key]
+    except KeyError:
+        pass
 
     # Need to define the mode for the first time
     if upper_string == "MODE":
@@ -577,7 +548,7 @@ def get_mode(orig_string):
     if config.optimizer_requiring:
         ret = ret.requiring(*config.optimizer_requiring.split(":"))
     # Cache the mode for next time
-    _CACHED_RUNTIME_MODES[upper_string] = ret
+    _CACHED_RUNTIME_MODES[cache_key] = ret
 
     return ret
 

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/compile/profiling.py

@@ -1464,7 +1464,6 @@ class ProfileStats:
             ps.GE,
             ps.EQ,
             ps.NEQ,
-            ps.InRange,
             ps.Switch,
             ps.OR,
             ps.XOR,

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/configdefaults.py

@@ -45,8 +45,12 @@ def _filter_mode(val):
         "NanGuardMode",
         "FAST_COMPILE",
         "DEBUG_MODE",
+        "CVM",
+        "C",
         "JAX",
         "NUMBA",
+        "PYTORCH",
+        "MLX",
     ]
     if val in str_options:
         return val
@@ -367,13 +371,26 @@ def add_compile_configvars():
     )
     del param
 
+    default_linker = "auto"
+
     if rc == 0 and config.cxx != "":
         # Keep the default linker the same as the one for the mode FAST_RUN
-        linker_options = ["c|py", "py", "c", "c|py_nogc", "vm", "vm_nogc", "cvm_nogc"]
+        linker_options = [
+            "cvm",
+            "c|py",
+            "py",
+            "c",
+            "c|py_nogc",
+            "vm",
+            "vm_nogc",
+            "cvm_nogc",
+            "numba",
+            "jax",
+        ]
     else:
         # g++ is not present or the user disabled it,
         # linker should default to python only.
-        linker_options = ["py", "vm_nogc"]
+        linker_options = ["py", "vm", "vm_nogc", "numba", "jax"]
         if type(config).cxx.is_default:
             # If the user provided an empty value for cxx, do not warn.
             _logger.warning(
@@ -385,9 +402,8 @@ def add_compile_configvars():
 
     config.add(
         "linker",
-        "Default linker used if the pytensor flags mode is Mode",
-        # Not mutable because the default mode is cached after the first use.
-        EnumStr("cvm", linker_options, mutable=False),
+        "Default linker used if the pytensor flags mode is Mode or FAST_RUN",
+        EnumStr(default_linker, linker_options, mutable=True),
         in_c_key=False,
     )
 

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/configparser.py

@@ -410,7 +410,9 @@ class ConfigParam:
                 f"The config parameter '{self.name}' was registered on a different instance of the PyTensorConfigParser."
                 f" It is not accessible through the instance with id '{id(cls)}' because of safeguarding."
             )
-        if not hasattr(self, "val"):
+        try:
+            return self.val
+        except AttributeError:
             try:
                 val_str = cls.fetch_val_for_key(self.name, delete_key=delete_key)
                 self.is_default = False

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/gradient.py

@@ -1784,14 +1784,14 @@ class numeric_grad:
 
 def mode_not_slow(mode):
     from pytensor.compile.debugmode import DebugMode
-    from pytensor.compile.mode import FAST_RUN, get_mode
+    from pytensor.compile.mode import get_mode
 
     if mode == "FAST_COMPILE":
-        return FAST_RUN
+        return get_mode("FAST_RUN")
     mode = get_mode(mode)
     if isinstance(mode, DebugMode):
         opt = mode.optimizer
-        return FAST_RUN.clone(optimizer=opt)
+        return get_mode("FAST_RUN").clone(optimizer=opt)
     else:
         return mode
 
@@ -2188,7 +2188,7 @@ def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):
         # It is possible that the inputs are disconnected from expr,
         # even if they are connected to cost.
         # This should not be an error.
-        hess, updates = pytensor.scan(
+        hess = pytensor.scan(
             lambda i, y, x: grad(
                 y[i],
                 x,
@@ -2197,9 +2197,7 @@ def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):
             ),
             sequences=pytensor.tensor.arange(expr.shape[0]),
             non_sequences=[expr, input],
-        )
-        assert not updates, (
-            "Scan has returned a list of updates; this should not happen."
+            return_updates=False,
         )
         hessians.append(hess)
     return as_list_or_tuple(using_list, using_tuple, hessians)

{pytensor-2.35.1 → pytensor-2.36.0}/pytensor/graph/rewriting/basic.py

@@ -1010,7 +1010,7 @@ class FromFunctionNodeRewriter(NodeRewriter):
 
 
 def node_rewriter(
-    tracks: Sequence[Op | type, OpPattern] | None,
+    tracks: Sequence[Op | type | OpPattern] | None,
     inplace: bool = False,
     requirements: tuple[type, ...] | None = (),
 ):