pytensor 3.0.3__tar.gz → 3.0.4__tar.gz

{pytensor-3.0.3/pytensor.egg-info → pytensor-3.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.0.3
+Version: 3.0.4
 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.0.3 → pytensor-3.0.4}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-05-20T17:31:39-0500",
+ "date": "2026-06-02T00:24:22+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "b712ea6e6a7541ce24c47f8869699178958dadf6",
- "version": "3.0.3"
+ "full-revisionid": "1df3475ac881a6bb182828c3dc8e36ed0439827f",
+ "version": "3.0.4"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/assumptions/__init__.py

@@ -6,8 +6,10 @@ import pytensor.assumptions.blockwise
 import pytensor.assumptions.diagonal
 import pytensor.assumptions.dimshuffle
 import pytensor.assumptions.orthogonal
+import pytensor.assumptions.permutation
 import pytensor.assumptions.positive_definite
 import pytensor.assumptions.reshape
+import pytensor.assumptions.selection
 import pytensor.assumptions.shape
 import pytensor.assumptions.subtensor
 import pytensor.assumptions.symmetric
@@ -18,7 +20,9 @@ from pytensor.assumptions.core import (
     IMPLIES,
     LOWER_TRIANGULAR,
     ORTHOGONAL,
+    PERMUTATION,
     POSITIVE_DEFINITE,
+    SELECTION,
     SYMMETRIC,
     UPPER_TRIANGULAR,
     AssumptionFeature,

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/assumptions/core.py

@@ -101,6 +101,8 @@ UPPER_TRIANGULAR = AssumptionKey("upper_triangular", short_name="triu")
 SYMMETRIC = AssumptionKey("symmetric", short_name="sym")
 POSITIVE_DEFINITE = AssumptionKey("positive_definite", short_name="pd")
 ORTHOGONAL = AssumptionKey("orthogonal", short_name="orth")
+SELECTION = AssumptionKey("selection", short_name="sel")
+PERMUTATION = AssumptionKey("permutation", short_name="perm")
 
 ALL_KEYS = (
     DIAGONAL,
@@ -109,11 +111,14 @@ ALL_KEYS = (
     SYMMETRIC,
     POSITIVE_DEFINITE,
     ORTHOGONAL,
+    SELECTION,
+    PERMUTATION,
 )
 
 # Implications about structural properties derivably from other structural properties
 register_implies(DIAGONAL, LOWER_TRIANGULAR, UPPER_TRIANGULAR, SYMMETRIC)
 register_implies(POSITIVE_DEFINITE, SYMMETRIC)
+register_implies(PERMUTATION, SELECTION, ORTHOGONAL)
 
 
 def register_assumption(
@@ -211,6 +216,9 @@ class AssumptionFeature(Feature):
             else:
                 self.cache[(new_var, key)] = FactState.join(old_state, new_state)
             self._var_to_keys.setdefault(new_var, set()).add(key)
+            # new_var gained a fact whose implications were cached UNKNOWN before it was
+            # known; flag it so the next get() sweep re-derives those stale entries.
+            self._stale_vars.add(new_var)
 
         self._stale_vars.update(node.outputs)
 

pytensor-3.0.4/pytensor/assumptions/permutation.py

@@ -0,0 +1,75 @@
+import numpy as np
+
+from pytensor.assumptions.alloc import eye_identity_rule
+from pytensor.assumptions.core import (
+    PERMUTATION,
+    FactState,
+    all_inputs_have_key,
+    propagate_first,
+    register_assumption,
+    register_constant_inference,
+)
+from pytensor.assumptions.dimshuffle import left_expand_dims_propagates_matrix_property
+from pytensor.assumptions.subtensor import subtensor_propagates_matrix_property
+from pytensor.tensor.basic import Eye
+from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.linalg.constructors import BlockDiagonal
+from pytensor.tensor.linalg.inverse import MatrixInverse, MatrixPinv
+from pytensor.tensor.linalg.products import KroneckerProduct
+from pytensor.tensor.math import Dot
+from pytensor.tensor.subtensor import Subtensor
+from pytensor.tensor.variable import TensorConstant
+
+
+# A PERMUTATION matrix is a square 0/1 matrix with exactly one 1 in every row and every
+# column -- a reordering of the identity's columns. It is a square SELECTION and is
+# ORTHOGONAL; unlike a general selection it is closed under transpose and inverse
+# (P.T = P^{-1} is again a permutation).
+
+
+def _permutation_from_constant(var: TensorConstant) -> FactState:
+    """Recognize a constant PERMUTATION matrix from its data, caching the scan on
+    ``var.tag.is_permutation``."""
+    cached: FactState | None = getattr(var.tag, "is_permutation", None)
+    if cached is not None:
+        return cached
+
+    data = np.asarray(var.data)
+    if data.ndim < 2 or data.shape[-1] != data.shape[-2]:
+        result = FactState.FALSE
+    else:
+        # The row/column-sum reductions are cheaper and short-circuit the full-size binary
+        # scan; a doubly-stochastic matrix shows the binary check is still required.
+        is_permutation = (
+            bool(np.all(data.sum(axis=-2) == 1))
+            and bool(np.all(data.sum(axis=-1) == 1))
+            and bool(np.all((data == 0) | (data == 1)))
+        )
+        result = FactState.TRUE if is_permutation else FactState.FALSE
+
+    var.tag.is_permutation = result
+    return result
+
+
+register_constant_inference(PERMUTATION, _permutation_from_constant)
+
+# eye(n, m, k) is a permutation iff it is the square identity -- the same condition that
+# makes it orthogonal.
+register_assumption(PERMUTATION, Eye)(eye_identity_rule)
+# The symmetric group is closed under matmul and inversion.
+register_assumption(PERMUTATION, Dot)(all_inputs_have_key)
+register_assumption(PERMUTATION, MatrixInverse)(propagate_first)
+register_assumption(PERMUTATION, MatrixPinv)(propagate_first)
+register_assumption(PERMUTATION, BlockDiagonal)(all_inputs_have_key)
+register_assumption(PERMUTATION, KroneckerProduct)(all_inputs_have_key)
+register_assumption(PERMUTATION, Subtensor)(subtensor_propagates_matrix_property)
+
+
+@register_assumption(PERMUTATION, DimShuffle)
+def _dimshuffle(key, op, feature, fgraph, node, input_states):
+    """Permutation survives matrix transpose (P.T = P^{-1}) and batch left-expand-dims."""
+    if input_states[0] is not FactState.TRUE:
+        return [FactState.UNKNOWN]
+    if op.is_matrix_transpose or left_expand_dims_propagates_matrix_property(op):
+        return [FactState.TRUE]
+    return [FactState.UNKNOWN]

pytensor-3.0.4/pytensor/assumptions/selection.py

@@ -0,0 +1,127 @@
+import numpy as np
+
+from pytensor.assumptions.core import (
+    SELECTION,
+    FactState,
+    all_inputs_have_key,
+    register_assumption,
+    register_constant_inference,
+    true_if,
+)
+from pytensor.assumptions.subtensor import subtensor_propagates_matrix_property
+from pytensor.graph.basic import Variable
+from pytensor.tensor.basic import Eye
+from pytensor.tensor.linalg.constructors import BlockDiagonal
+from pytensor.tensor.linalg.products import KroneckerProduct
+from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor
+from pytensor.tensor.variable import TensorConstant
+
+
+def _selection_from_constant(var: TensorConstant) -> FactState:
+    """Recognize a constant SELECTION matrix from its data, caching the O(n*k) scan on
+    ``var.tag.is_selection``."""
+    cached: FactState | None = getattr(var.tag, "is_selection", None)
+    if cached is not None:
+        return cached
+
+    data = np.asarray(var.data)
+    if data.ndim < 2:
+        result = FactState.FALSE
+    else:
+        if not (data.sum(axis=-2) == 1).all():
+            is_selection = False
+        elif data.dtype.kind in "uib":
+            is_selection = data.max(initial=1) <= 1 and data.min(initial=0) >= 0
+        else:
+            is_selection = bool(((data == 0) | (data == 1)).all())
+        result = FactState.TRUE if is_selection else FactState.FALSE
+
+    var.tag.is_selection = result
+    return result
+
+
+register_constant_inference(SELECTION, _selection_from_constant)
+
+
+def column_selection_index(op, node) -> Variable | None:
+    """Return the 1-D integer index of an ``x[..., idx]`` advanced subtensor, or ``None``.
+
+    Matches every axis but the last a full slice, the last an integer vector. Selecting
+    columns preserves selection for any ``idx``; selecting rows does not, so only this
+    pattern is recognized.
+    """
+    x = node.inputs[0]
+    ndim = x.type.ndim
+    idx_list = op.idx_list
+    if ndim < 2 or len(idx_list) != ndim:
+        return None
+
+    full = slice(None, None, None)
+    if any(entry != full for entry in idx_list[:-1]):
+        return None
+
+    last = idx_list[-1]
+    if not isinstance(last, int) or isinstance(last, bool):
+        return None
+
+    index_var: Variable = node.inputs[1:][last]
+    if index_var.type.ndim != 1 or not index_var.type.dtype.startswith(("int", "uint")):
+        return None
+    return index_var
+
+
+@register_assumption(SELECTION, Eye)
+def _eye(key, op, feature, fgraph, node, input_states):
+    """``eye(n, m, k)`` is a selection iff ``k == 0`` (main diagonal) and ``n >= m``
+    (no trailing zero columns)."""
+    n, m, k = node.inputs
+    if not isinstance(k, TensorConstant):
+        return [FactState.UNKNOWN]
+    if k.data.item() != 0:
+        return [FactState.FALSE]
+    if n is m:
+        return [FactState.TRUE]
+    if isinstance(n, TensorConstant) and isinstance(m, TensorConstant):
+        return true_if(n.data.item() >= m.data.item(), else_false=True)
+    rows, cols = node.outputs[0].type.shape[-2:]
+    if rows is not None and cols is not None:
+        return true_if(rows >= cols, else_false=True)
+    return [FactState.UNKNOWN]
+
+
+register_assumption(SELECTION, BlockDiagonal)(all_inputs_have_key)
+register_assumption(SELECTION, KroneckerProduct)(all_inputs_have_key)
+
+
+@register_assumption(SELECTION, AdvancedSubtensor)
+def _advanced_subtensor(key, op, feature, fgraph, node, input_states):
+    if (
+        input_states[0] is FactState.TRUE
+        and column_selection_index(op, node) is not None
+    ):
+        return [FactState.TRUE]
+    return [FactState.UNKNOWN]
+
+
+@register_assumption(SELECTION, Subtensor)
+def _subtensor(key, op, feature, fgraph, node, input_states):
+    # Batch-axis indexing (trailing two axes untouched) keeps every per-matrix property.
+    states = subtensor_propagates_matrix_property(
+        key, op, feature, fgraph, node, input_states
+    )
+    if states[0] is FactState.TRUE:
+        return states
+
+    # A column slice keeps a subset of the one-hot columns -- still a selection.
+    if input_states[0] is FactState.TRUE:
+        x = node.inputs[0]
+        idx_list = op.idx_list
+        full = slice(None, None, None)
+        if (
+            2 <= x.type.ndim == len(idx_list)
+            and all(entry == full for entry in idx_list[:-1])
+            and isinstance(idx_list[-1], slice)
+        ):
+            return [FactState.TRUE]
+
+    return [FactState.UNKNOWN]

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/assumptions/specify.py

@@ -67,6 +67,8 @@ def assume(
     symmetric: bool | None = None,
     positive_definite: bool | None = None,
     orthogonal: bool | None = None,
+    selection: bool | None = None,
+    permutation: bool | None = None,
 ):
     """Attach structural assumptions to a symbolic tensor.
 
@@ -90,6 +92,10 @@ def assume(
         Assert that *x* is (or is not) positive-definite.
     orthogonal : bool, optional
         Assert that *x* is (or is not) orthogonal.
+    selection : bool, optional
+        Assert that *x* is (or is not) a selection matrix
+    permutation : bool, optional
+        Assert that *x* is (or is not) a permutation matrix
 
     Returns
     -------
@@ -113,6 +119,8 @@ def assume(
         "symmetric": symmetric,
         "positive_definite": positive_definite,
         "orthogonal": orthogonal,
+        "selection": selection,
+        "permutation": permutation,
     }
     assumptions = {
         name: FactState.TRUE if value else FactState.FALSE

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/assumptions/subtensor.py

@@ -1,6 +1,10 @@
 from pytensor.assumptions.core import (
     ALL_KEYS,
-    ORTHOGONAL,
+    DIAGONAL,
+    LOWER_TRIANGULAR,
+    POSITIVE_DEFINITE,
+    SYMMETRIC,
+    UPPER_TRIANGULAR,
     FactState,
     register_assumption,
     true_if,
@@ -46,7 +50,7 @@ def incsubtensor_propagates_matrix_property(
     - ``set`` over a partial batch region: every output slice comes from ``base``
       or ``value``, so both must carry the property.
     - ``inc``: the written region becomes ``base + value``; the property must be
-      closed under addition -- true for every key except ``orthogonal``.
+      closed under addition (diagonal, triangular, symmetric, positive-definite).
 
     An index that reaches into the core axes is left to the per-key rules (e.g.
     the diagonal-position write handled by ``indexes_diagonal``).
@@ -66,7 +70,14 @@ def incsubtensor_propagates_matrix_property(
             return [value_state]
         return true_if(base_state is FactState.TRUE and value_state is FactState.TRUE)
 
-    if key is ORTHOGONAL:
+    # Closed under addition
+    if key not in (
+        DIAGONAL,
+        LOWER_TRIANGULAR,
+        UPPER_TRIANGULAR,
+        SYMMETRIC,
+        POSITIVE_DEFINITE,
+    ):
         return [FactState.UNKNOWN]
     return true_if(base_state is FactState.TRUE and value_state is FactState.TRUE)
 

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/assumptions/triangular.py

@@ -136,21 +136,32 @@ def _lu_upper(key, op, feature, fgraph, node, input_states):
     return states
 
 
-def _dimshuffle_left_expand_dims(key, op, feature, fgraph, node, input_states):
-    """Triangularity survives left-expand-dims (batch broadcast).
+@register_assumption(LOWER_TRIANGULAR, DimShuffle)
+def _lower_dimshuffle(key, op, feature, fgraph, node, input_states):
+    """Mirror lower-triangularity through left-expand-dims and matrix transpose.
 
-    Matrix transpose swaps lower<->upper, so it is *not* propagated under
-    the same key.
+    Left-expand-dims preserves the matrix exactly and transpose swaps the two
+    triangles, swapping their fact states.
     """
-    if input_states[0] is not FactState.TRUE:
-        return [FactState.UNKNOWN]
+    if op.is_matrix_transpose:
+        return [feature.get(node.inputs[0], UPPER_TRIANGULAR)]
     if left_expand_dims_propagates_matrix_property(op):
-        return [FactState.TRUE]
+        return [input_states[0]]
     return [FactState.UNKNOWN]
 
 
-register_assumption(LOWER_TRIANGULAR, DimShuffle)(_dimshuffle_left_expand_dims)
-register_assumption(UPPER_TRIANGULAR, DimShuffle)(_dimshuffle_left_expand_dims)
+@register_assumption(UPPER_TRIANGULAR, DimShuffle)
+def _upper_dimshuffle(key, op, feature, fgraph, node, input_states):
+    """Mirror upper-triangularity through left-expand-dims and matrix transpose.
+
+    Left-expand-dims preserves the matrix exactly and transpose swaps the two
+    triangles, swapping their fact states.
+    """
+    if op.is_matrix_transpose:
+        return [feature.get(node.inputs[0], LOWER_TRIANGULAR)]
+    if left_expand_dims_propagates_matrix_property(op):
+        return [input_states[0]]
+    return [FactState.UNKNOWN]
 
 
 @register_assumption(LOWER_TRIANGULAR, Elemwise)

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/compile/executor.py

@@ -395,6 +395,23 @@ class Function:
         )
         fg_cpy.update_mapping = maker.fgraph.update_mapping
 
+        # Warn if any shared variable with a deleted update is being
+        # destroyed (mutated inplace) by a node in the copied graph.
+        # In that case the shared variable storage will still be mutated
+        # even though the update is not applied.
+        if delete_updates and hasattr(maker.fgraph, "destroyers"):
+            for in_idx in maker.fgraph.update_mapping.values():
+                input_var = maker.fgraph.inputs[in_idx]
+                for destroyer in maker.fgraph.destroyers(input_var):
+                    if destroyer in memo:
+                        warnings.warn(
+                            f"Shared variable '{input_var.name}' will still be mutated "
+                            "even though its update is being deleted, because an "
+                            "operation in the graph modifies it inplace.",
+                            UserWarning,
+                            stacklevel=2,
+                        )
+
         # Re initialize Outs and swap update and variable in Ins
         # By doing this, we can pass FunctionMaker.check_unused_inputs()
         if delete_updates:

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/graph/traversal.py

@@ -526,9 +526,11 @@ def walk_toposort(
 
     Notes
     -----
-
     ``deps(i)`` should behave like a pure function (no funny business with internal state).
 
+    The implementation uses identity (``is``), so items returned by ``deps``
+    must be the *same objects* as those yielded by ``graphs``.
+
     The order of the return value list is determined by the order of nodes
     returned by the `deps` function.
     """

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/ifelse.py

@@ -423,21 +423,16 @@ def ifelse(
 @node_rewriter([IfElse])
 def cond_make_inplace(fgraph, node):
     op = node.op
-    if (
-        isinstance(op, IfElse)
-        and not op.as_view
-        and
-        # For big graph, do not make inplace scalar to speed up
-        # optimization.
-        (
-            len(fgraph.apply_nodes) < 500
-            or not all(getattr(o.type, "ndim", -1) == 0 for o in node.outputs)
-        )
+    if op.as_view:
+        return False
+    # For big graph, do not make inplace scalar to speed up optimization.
+    if len(fgraph.apply_nodes) >= 500 and all(
+        getattr(o.type, "ndim", -1) == 0 for o in node.outputs
     ):
-        return IfElse(n_outs=op.n_outs, as_view=True, name=op.name)(
-            *node.inputs, return_list=True
-        )
-    return False
+        return False
+    return IfElse(n_outs=op.n_outs, as_view=True, name=op.name)(
+        *node.inputs, return_list=True
+    )
 
 
 optdb.register(
@@ -549,8 +544,6 @@ def ifelse_lift_single_if_through_acceptable_ops(fgraph, main_node):
 @node_rewriter([IfElse])
 def cond_merge_ifs_true(fgraph, node):
     op = node.op
-    if not isinstance(op, IfElse):
-        return False
     t_ins = node.inputs[1:][: op.n_outs]
 
     replace = {}
@@ -576,8 +569,6 @@ def cond_merge_ifs_true(fgraph, node):
 @node_rewriter([IfElse])
 def cond_merge_ifs_false(fgraph, node):
     op = node.op
-    if not isinstance(op, IfElse):
-        return False
     f_ins = node.inputs[1:][op.n_outs :]
 
     replace = {}
@@ -656,8 +647,6 @@ class CondMerge(GraphRewriter):
 def cond_remove_identical(fgraph, node):
     op = node.op
 
-    if not isinstance(op, IfElse):
-        return False
     aes = node.inputs[1:][: op.n_outs]
     fs = node.inputs[1:][op.n_outs :]
 

{pytensor-3.0.3 → pytensor-3.0.4}/pytensor/link/numba/dispatch/blockwise.py

@@ -11,6 +11,9 @@ from pytensor.link.numba.dispatch.basic import (
     register_funcify_and_cache_key,
 )
 from pytensor.link.numba.dispatch.vectorize_codegen import (
+    NO_INDEXED_INPUTS,
+    NO_INDEXED_OUTPUTS,
+    NO_SIZE,
     _jit_options,
     _vectorized,
     encode_literals,
@@ -90,7 +93,9 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
                 (),  # constant_inputs
                 inputs,
                 tuple_core_shapes,
-                None,  # size
+                NO_SIZE,
+                NO_INDEXED_INPUTS,
+                NO_INDEXED_OUTPUTS,
             )
 
         return impl