pytensor 2.28.1__tar.gz → 2.28.3__tar.gz

{pytensor-2.28.1/pytensor.egg-info → pytensor-2.28.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytensor
-Version: 2.28.1
+Version: 2.28.3
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-02-24T18:22:58+0100",
+ "date": "2025-03-05T13:02:56+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "67017a662a65abaa0e388925762aa3c903abe840",
- "version": "2.28.1"
+ "full-revisionid": "89d5366c888f430003aa6f16a13ad9cdd2826d95",
+ "version": "2.28.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/compile/debugmode.py

@@ -1966,6 +1966,12 @@ class _Maker(FunctionMaker):  # inheritance buys a few helper functions
         If the outputs argument for pytensor.function was a list, then
         output_keys is None. If the outputs argument was a dict, then
         output_keys is a sorted list of the keys from that dict.
+    trust_input : bool, default False
+        If True, no input validation checks are performed when the function is
+        called. This includes checking the number of inputs, their types and
+        that multiple inputs are not aliased to each other. Failure to meet any
+        of these conditions can lead to computational errors or to the
+        interpreter crashing.
 
     Notes
     -----
@@ -1993,6 +1999,7 @@ class _Maker(FunctionMaker):  # inheritance buys a few helper functions
         output_keys=None,
         name=None,
         no_fgraph_prep=False,
+        trust_input=False,
     ):
         self.mode = mode
         self.profile = profile
@@ -2146,6 +2153,7 @@ class _Maker(FunctionMaker):  # inheritance buys a few helper functions
         self.on_unused_input = on_unused_input  # Used for the pickling/copy
         self.output_keys = output_keys
         self.name = name
+        self.trust_input = trust_input
 
         self.required = [(i.value is None) for i in self.inputs]
         self.refeed = [

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/compile/function/__init__.py

@@ -37,6 +37,7 @@ def function_dump(
     profile: bool | ProfileStats | None = None,
     on_unused_input: str | None = None,
     extra_tag_to_remove: str | None = None,
+    trust_input: bool = False,
 ):
     """
     This is helpful to make a reproducible case for problems during PyTensor
@@ -82,6 +83,7 @@ def function_dump(
         "allow_input_downcast": allow_input_downcast,
         "profile": profile,
         "on_unused_input": on_unused_input,
+        "trust_input": trust_input,
     }
     with Path(filename).open("wb") as f:
         pickler = pytensor.misc.pkl_utils.StripPickler(
@@ -107,6 +109,7 @@ def function(
     allow_input_downcast: bool | None = None,
     profile: bool | ProfileStats | None = None,
     on_unused_input: str | None = None,
+    trust_input: bool = False,
 ):
     """
     Return a :class:`callable object <pytensor.compile.function.types.Function>`
@@ -164,6 +167,12 @@ def function(
     on_unused_input
         What to do if a variable in the 'inputs' list is not used in the graph.
         Possible values are 'raise', 'warn', 'ignore' and None.
+    trust_input: bool, default False
+        If True, no input validation checks are performed when the function is
+        called. This includes checking the number of inputs, their types and
+        that multiple inputs are not aliased to each other. Failure to meet any
+        of these conditions can lead to computational errors or to the
+        interpreter crashing.
 
     Returns
     -------
@@ -310,7 +319,12 @@ def function(
                 "semantics, which disallow using updates and givens"
             )
         fn = orig_function(
-            inputs, outputs, mode=mode, accept_inplace=accept_inplace, name=name
+            inputs,
+            outputs,
+            mode=mode,
+            accept_inplace=accept_inplace,
+            name=name,
+            trust_input=trust_input,
         )
     else:
         # note: pfunc will also call orig_function -- orig_function is
@@ -329,5 +343,6 @@ def function(
             on_unused_input=on_unused_input,
             profile=profile,
             output_keys=output_keys,
+            trust_input=trust_input,
         )
     return fn

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/compile/function/pfunc.py

@@ -377,6 +377,7 @@ def pfunc(
     on_unused_input=None,
     output_keys=None,
     fgraph: FunctionGraph | None = None,
+    trust_input: bool = False,
 ) -> Function:
     """
     Function-constructor for graphs with shared variables.
@@ -425,6 +426,12 @@ def pfunc(
     fgraph
         An existing `FunctionGraph` from which to construct the returned
         `Function`.  When this is non-``None``, nothing is cloned.
+    trust_input : bool, default False
+        If True, no input validation checks are performed when the function is
+        called. This includes checking the number of inputs, their types and
+        that multiple inputs are not aliased to each other. Failure to meet any
+        of these conditions can lead to computational errors or to the
+        interpreter crashing.
 
     Returns
     -------
@@ -472,6 +479,7 @@ def pfunc(
         on_unused_input=on_unused_input,
         output_keys=output_keys,
         fgraph=fgraph,
+        trust_input=trust_input,
     )
 
 

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/compile/function/types.py

@@ -5,6 +5,7 @@ import copyreg
 import logging
 import time
 import warnings
+from collections.abc import Sequence
 from itertools import chain
 from typing import TYPE_CHECKING
 
@@ -168,6 +169,59 @@ class Supervisor(Feature):
                 raise InconsistencyError(f"Trying to destroy a protected variable: {r}")
 
 
+def add_supervisor_to_fgraph(
+    fgraph: FunctionGraph,
+    input_specs: Sequence[SymbolicInput],
+    accept_inplace: bool = False,
+) -> None:
+    """Setup Supervisor Feature in a FunctionGraph, so that inplace rewrites can be used.
+
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        The FunctionGraph to setup the Supervisor Feature in.
+    input_specs: Sequence of SymbolicInput
+        The input specifications for the FunctionGraph.
+        Inputs with the attribute `mutable=False` and which are not already destroyed by an inplace operation
+        (if `accept_inplace` is True) will be protected from inplace operations.
+        Otherwise, they will be allowed to be destroyed.
+    accept_inplace: bool
+        Whether to allow inplace operations to already be present in the graph.
+
+    Raises
+    ------
+    TypeError
+        If inplace operations are not allowed and the graph already contains inplace operations.
+
+    """
+
+    has_destroy_handler = hasattr(fgraph, "destroyers")
+    if not (has_destroy_handler and accept_inplace):
+        # Check if fgraph already contains destructive operations,
+        # in which case we need to add a DestroyHandler or raise an error
+        for node in fgraph.apply_nodes:
+            if node.op.destroy_map:
+                if not accept_inplace:
+                    raise TypeError(
+                        f"Graph must not contain inplace operations: {node}"
+                    )
+                else:
+                    has_destroy_handler = True
+                    fgraph.attach_feature(DestroyHandler())
+                    break
+
+    # Protect all immutable inputs from inplace operations.
+    fgraph.attach_feature(
+        Supervisor(
+            input
+            for spec, input in zip(input_specs, fgraph.inputs, strict=True)
+            if not (
+                spec.mutable or has_destroy_handler and fgraph.has_destroyers([input])
+            )
+        )
+    )
+
+
 def std_fgraph(
     input_specs: list[SymbolicInput],
     output_specs: list[SymbolicOutput],
@@ -229,24 +283,8 @@ def std_fgraph(
 
         found_updates.extend(map(SymbolicOutput, updates))
 
-    for node in fgraph.apply_nodes:
-        if node.op.destroy_map:
-            if not accept_inplace:
-                raise TypeError(f"Graph must not contain inplace operations: {node}")
-            else:
-                fgraph.attach_feature(DestroyHandler())
-                break
-
-    # We need to protect all immutable inputs from inplace operations.
-    fgraph.attach_feature(
-        Supervisor(
-            input
-            for spec, input in zip(input_specs, fgraph.inputs, strict=True)
-            if not (
-                spec.mutable
-                or (hasattr(fgraph, "destroyers") and fgraph.has_destroyers([input]))
-            )
-        )
+    add_supervisor_to_fgraph(
+        fgraph=fgraph, input_specs=input_specs, accept_inplace=accept_inplace
     )
 
     # If named nodes are replaced, keep the name
@@ -335,6 +373,7 @@ class Function:
         return_none: bool,
         output_keys,
         maker: "FunctionMaker",
+        trust_input: bool = False,
         name: str | None = None,
     ):
         """
@@ -369,6 +408,12 @@ class Function:
             TODO
         maker
             The `FunctionMaker` that created this instance.
+        trust_input : bool, default False
+            If True, no input validation checks are performed when the function is
+            called. This includes checking the number of inputs, their types and
+            that multiple inputs are not aliased to each other. Failure to meet any
+            of these conditions can lead to computational errors or to the
+            interpreter crashing.
         name
             A string name.
         """
@@ -382,7 +427,7 @@ class Function:
         self.return_none = return_none
         self.maker = maker
         self.profile = None  # reassigned in FunctionMaker.create
-        self.trust_input = False  # If True, we don't check the input parameter
+        self.trust_input = trust_input  # If True, we don't check the input parameter
         self.name = name
         self.nodes_with_inner_function = []
         self.output_keys = output_keys
@@ -1303,7 +1348,12 @@ class FunctionMaker:
     name : str
         An optional name for this function. If used, the profile mode will
         print the time spent in this function.
-
+    trust_input : bool, default False
+        If True, no input validation checks are performed when the function is
+        called. This includes checking the number of inputs, their types and
+        that multiple inputs are not aliased to each other. Failure to meet any
+        of these conditions can lead to computational errors or to the
+        interpreter crashing.
     """
 
     @staticmethod
@@ -1469,6 +1519,7 @@ class FunctionMaker:
         output_keys=None,
         name=None,
         no_fgraph_prep=False,
+        trust_input=False,
     ):
         # Save the provided mode, not the instantiated mode.
         # The instantiated mode don't pickle and if we unpickle an PyTensor
@@ -1571,6 +1622,7 @@ class FunctionMaker:
         self.on_unused_input = on_unused_input  # Used for the pickling/copy
         self.output_keys = output_keys
         self.name = name
+        self.trust_input = trust_input
 
         self.required = [(i.value is None) for i in self.inputs]
         self.refeed = [
@@ -1688,6 +1740,7 @@ class FunctionMaker:
             self.return_none,
             self.output_keys,
             self,
+            trust_input=self.trust_input,
             name=self.name,
         )
 
@@ -1705,6 +1758,7 @@ def orig_function(
     on_unused_input=None,
     output_keys=None,
     fgraph: FunctionGraph | None = None,
+    trust_input: bool = False,
 ) -> Function:
     """
     Return a Function that will calculate the outputs from the inputs.
@@ -1735,7 +1789,12 @@ def orig_function(
     fgraph
         An existing `FunctionGraph` to use instead of constructing a new one
         from cloned `outputs`.
-
+    trust_input : bool, default False
+        If True, no input validation checks are performed when the function is
+        called. This includes checking the number of inputs, their types and
+        that multiple inputs are not aliased to each other. Failure to meet any
+        of these conditions can lead to computational errors or to the
+        interpreter crashing.
     """
 
     if profile:
@@ -1768,6 +1827,7 @@ def orig_function(
             output_keys=output_keys,
             name=name,
             fgraph=fgraph,
+            trust_input=trust_input,
         )
         with config.change_flags(compute_test_value="off"):
             fn = m.create(defaults)

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/compile/mode.py

@@ -138,7 +138,11 @@ class AddDestroyHandler(GraphRewriter):
                 break
         if not supervisor_added:
             warnings.warn(
-                f"A Supervisor feature is missing from {fgraph}.",
+                (
+                    f"A Supervisor feature is missing from {fgraph}.\n"
+                    "This is needed for inplace rewrites. Either exclude inplace rewrites or add a Supervisor feature.\n"
+                    "A Supervisor feature can be added via `pytensor.compile.function.types.add_supervisor_to_fgraph`."
+                ),
                 stacklevel=3,
             )
 

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/configdefaults.py

@@ -418,11 +418,18 @@ def add_compile_configvars():
 
     config.add(
         "optimizer_verbose",
-        "If True, we print all optimization being applied",
+        "Print information about rewrites that are applied during a graph transformation.",
         BoolParam(False),
         in_c_key=False,
     )
 
+    config.add(
+        "optimizer_verbose_ignore",
+        "Do not print information for rewrites with these names when `optimizer_verbose` is `True`. Separate names with ','",
+        StrParam(""),
+        in_c_key=False,
+    )
+
     config.add(
         "on_opt_error",
         (

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/configparser.py

@@ -81,6 +81,7 @@ class PyTensorConfigParser:
     allow_gc: bool
     optimizer: str
     optimizer_verbose: bool
+    optimizer_verbose_ignore: str
     on_opt_error: str
     nocleanup: bool
     on_unused_input: str

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/graph/features.py

@@ -567,6 +567,13 @@ class ReplaceValidate(History, Validator):
         if verbose is None:
             verbose = config.optimizer_verbose
 
+        if verbose:
+            print_reason = True
+            if config.optimizer_verbose_ignore:
+                print_reason = str(reason) not in config.optimizer_verbose_ignore.split(
+                    ","
+                )
+
         for r, new_r in replacements:
             try:
                 fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
@@ -608,7 +615,7 @@ class ReplaceValidate(History, Validator):
                 )
             raise
 
-        if verbose:
+        if verbose and print_reason:
             print(  # noqa: T201
                 f"rewriting: rewrite {reason} replaces {r} of {r.owner} with {new_r} of {new_r.owner}"
             )

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/graph/fg.py

@@ -490,10 +490,18 @@ class FunctionGraph(MetaObject):
         """
         if verbose is None:
             verbose = config.optimizer_verbose
+
         if verbose:
-            print(  # noqa: T201
-                f"rewriting: rewrite {reason} replaces {var} of {var.owner} with {new_var} of {new_var.owner}"
-            )
+            print_reason = True
+            if config.optimizer_verbose_ignore:
+                print_reason = str(reason) not in config.optimizer_verbose_ignore.split(
+                    ","
+                )
+
+            if print_reason:
+                print(  # noqa: T201
+                    f"rewriting: rewrite {reason} replaces {var} of {var.owner} with {new_var} of {new_var.owner}"
+                )
 
         new_var = var.type.filter_variable(new_var, allow_convert=True)
 

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/graph/rewriting/basic.py

@@ -1305,9 +1305,15 @@ class SequentialNodeRewriter(NodeRewriter):
                     new_vars = list(new_repl.values())
 
                 if config.optimizer_verbose:
-                    print(  # noqa: T201
-                        f"rewriting: rewrite {rewrite} replaces node {node} with {new_repl}"
-                    )
+                    print_reason = True
+                    if config.optimizer_verbose_ignore:
+                        print_reason = str(
+                            rewrite
+                        ) not in config.optimizer_verbose_ignore.split(",")
+                    if print_reason:
+                        print(  # noqa: T201
+                            f"rewriting: rewrite {rewrite} replaces node {node} with {new_repl}"
+                        )
 
                 if self.profile:
                     self.node_created[rewrite] += len(

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/link/jax/dispatch/elemwise.py

@@ -79,12 +79,7 @@ def jax_funcify_DimShuffle(op, **kwargs):
         for augm in op.augment:
             shape.insert(augm, 1)
 
-        res = jnp.reshape(res, shape)
-
-        if not op.inplace:
-            res = jnp.copy(res)
-
-        return res
+        return jnp.reshape(res, shape)
 
     return dimshuffle
 

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/link/jax/dispatch/scan.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 
-from pytensor.compile.mode import JAX
+from pytensor.compile.mode import JAX, get_mode
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.scan.op import Scan
 
@@ -19,7 +19,9 @@ def jax_funcify_Scan(op: Scan, **kwargs):
         )
 
     # Optimize inner graph (exclude any defalut rewrites that are incompatible with JAX mode)
-    rewriter = op.mode_instance.excluding(*JAX._optimizer.exclude).optimizer
+    rewriter = (
+        get_mode(op.mode).including("jax").excluding(*JAX._optimizer.exclude).optimizer
+    )
     rewriter(op.fgraph)
     scan_inner_func = jax_funcify(op.fgraph, **kwargs)
 

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/link/jax/dispatch/slinalg.py

@@ -53,7 +53,6 @@ def jax_funcify_Solve(op, **kwargs):
 @jax_funcify.register(SolveTriangular)
 def jax_funcify_SolveTriangular(op, **kwargs):
     lower = op.lower
-    trans = op.trans
     unit_diagonal = op.unit_diagonal
     check_finite = op.check_finite
 
@@ -62,7 +61,7 @@ def jax_funcify_SolveTriangular(op, **kwargs):
             A,
             b,
             lower=lower,
-            trans=trans,
+            trans=0,  # this is handled by explicitly transposing A, so it will always be 0 when we get to here.
             unit_diagonal=unit_diagonal,
             check_finite=check_finite,
         )

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/link/numba/dispatch/basic.py

@@ -16,9 +16,10 @@ from numba.core.errors import NumbaWarning, TypingError
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
 from numba.extending import box, overload
 
-from pytensor import config
+from pytensor import In, config
 from pytensor.compile import NUMBA
 from pytensor.compile.builders import OpFromGraph
+from pytensor.compile.function.types import add_supervisor_to_fgraph
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
@@ -430,7 +431,13 @@ def numba_funcify_OpFromGraph(op, node=None, **kwargs):
     # TODO: Not sure this is the right place to do this, should we have a rewrite that
     #  explicitly triggers the optimization of the inner graphs of OpFromGraph?
     #  The C-code defers it to the make_thunk phase
-    NUMBA.optimizer(op.fgraph)
+    fgraph = op.fgraph
+    add_supervisor_to_fgraph(
+        fgraph=fgraph,
+        input_specs=[In(x, borrow=True, mutable=False) for x in fgraph.inputs],
+        accept_inplace=True,
+    )
+    NUMBA.optimizer(fgraph)
     fgraph_fn = numba_njit(numba_funcify(op.fgraph, **kwargs))
 
     if len(op.fgraph.outputs) == 1:

{pytensor-2.28.1 → pytensor-2.28.3}/pytensor/link/numba/dispatch/elemwise.py

@@ -4,6 +4,7 @@ from textwrap import dedent, indent
 import numba
 import numpy as np
 from numba.core.extending import overload
+from numpy.lib.stride_tricks import as_strided
 
 from pytensor.graph.op import Op
 from pytensor.link.numba.dispatch import basic as numba_basic
@@ -411,97 +412,38 @@ def numba_funcify_CAReduce(op, node, **kwargs):
 
 @numba_funcify.register(DimShuffle)
 def numba_funcify_DimShuffle(op, node, **kwargs):
-    shuffle = tuple(op.shuffle)
-    transposition = tuple(op.transposition)
-    augment = tuple(op.augment)
-    inplace = op.inplace
+    # We use `as_strided` to achieve the DimShuffle behavior of transposing and expanding/squezing dimensions in one call
+    # Numba doesn't currently support multiple expand/squeeze, and reshape is limited to contiguous arrays.
+    new_order = tuple(op._new_order)
+    shape_template = (1,) * node.outputs[0].ndim
+    strides_template = (0,) * node.outputs[0].ndim
 
-    ndim_new_shape = len(shuffle) + len(augment)
-
-    no_transpose = all(i == j for i, j in enumerate(transposition))
-    if no_transpose:
-
-        @numba_basic.numba_njit
-        def transpose(x):
-            return x
-
-    else:
+    if new_order == ():
+        # Special case needed because of https://github.com/numba/numba/issues/9933
 
         @numba_basic.numba_njit
-        def transpose(x):
-            return np.transpose(x, transposition)
-
-    shape_template = (1,) * ndim_new_shape
-
-    # When `len(shuffle) == 0`, the `shuffle_shape[j]` expression below
-    # is typed as `getitem(Tuple(), int)`, which has no implementation
-    # (since getting an item from an empty sequence doesn't make sense).
-    # To avoid this compile-time error, we omit the expression altogether.
-    if len(shuffle) > 0:
-        # Use the statically known shape if available
-        if all(length is not None for length in node.outputs[0].type.shape):
-            shape = node.outputs[0].type.shape
-
-            @numba_basic.numba_njit
-            def find_shape(array_shape):
-                return shape
-
-        else:
-
-            @numba_basic.numba_njit
-            def find_shape(array_shape):
-                shape = shape_template
-                j = 0
-                for i in range(ndim_new_shape):
-                    if i not in augment:
-                        length = array_shape[j]
-                        shape = numba_basic.tuple_setitem(shape, i, length)
-                        j = j + 1
-                return shape
-
-    else:
-
-        @numba_basic.numba_njit
-        def find_shape(array_shape):
-            return shape_template
-
-    if ndim_new_shape > 0:
-
-        @numba_basic.numba_njit
-        def dimshuffle_inner(x, shuffle):
-            x = transpose(x)
-            shuffle_shape = x.shape[: len(shuffle)]
-            new_shape = find_shape(shuffle_shape)
-
-            # FIXME: Numba's `array.reshape` only accepts C arrays.
-            res_reshape = np.reshape(np.ascontiguousarray(x), new_shape)
+        def squeeze_to_0d(x):
+            return as_strided(x, shape=(), strides=())
 
-            if not inplace:
-                return res_reshape.copy()
-            else:
-                return res_reshape
+        return squeeze_to_0d
 
     else:
 
         @numba_basic.numba_njit
-        def dimshuffle_inner(x, shuffle):
-            return np.reshape(np.ascontiguousarray(x), ())
-
-    # Without the following wrapper function we would see this error:
-    # E   No implementation of function Function(<built-in function getitem>) found for signature:
-    # E
-    # E    >>> getitem(UniTuple(int64 x 2), slice<a:b>)
-    # E
-    # E   There are 22 candidate implementations:
-    # E      - Of which 22 did not match due to:
-    # E      Overload of function 'getitem': File: <numerous>: Line N/A.
-    # E        With argument(s): '(UniTuple(int64 x 2), slice<a:b>)':
-    # E       No match.
-    # ...(on this line)...
-    # E           shuffle_shape = res.shape[: len(shuffle)]
-    @numba_basic.numba_njit(inline="always")
-    def dimshuffle(x):
-        return dimshuffle_inner(np.asarray(x), shuffle)
+        def dimshuffle(x):
+            old_shape = x.shape
+            old_strides = x.strides
+
+            new_shape = shape_template
+            new_strides = strides_template
+            for i, o in enumerate(new_order):
+                if o != -1:
+                    new_shape = numba_basic.tuple_setitem(new_shape, i, old_shape[o])
+                    new_strides = numba_basic.tuple_setitem(
+                        new_strides, i, old_strides[o]
+                    )
+
+            return as_strided(x, shape=new_shape, strides=new_strides)
 
     return dimshuffle