pytensor 3.0.4__tar.gz → 3.0.5__tar.gz

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

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-06-02T00:24:22+0200",
+ "date": "2026-06-11T03:57:58+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "1df3475ac881a6bb182828c3dc8e36ed0439827f",
- "version": "3.0.4"
+ "full-revisionid": "3411b300d7a120e7fb58b5ad5495391ddd74e189",
+ "version": "3.0.5"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/assumptions/diagonal.py

@@ -49,8 +49,13 @@ def _diagonal_from_constant(var: TensorConstant) -> FactState:
         if m != n:
             result = FactState.FALSE
         else:
-            eye_mask = np.eye(n, dtype=bool)
-            result = FactState.FALSE if np.any(data * ~eye_mask) else FactState.TRUE
+            # Off-diagonal has a nonzero iff there are more nonzeros in ``data``
+            # than on its diagonal.
+            diag = np.diagonal(data, axis1=-2, axis2=-1)
+            if np.count_nonzero(data) == np.count_nonzero(diag):
+                result = FactState.TRUE
+            else:
+                result = FactState.FALSE
 
     var.tag.is_diagonal = result
     return result

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/assumptions/permutation.py

@@ -38,13 +38,18 @@ def _permutation_from_constant(var: TensorConstant) -> FactState:
     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)))
-        )
+        with np.errstate(invalid="ignore"):
+            if not (data.sum(axis=-2) == 1).all():
+                is_permutation = False
+            elif not (data.sum(axis=-1) == 1).all():
+                is_permutation = False
+            elif data.dtype.kind in "ub":
+                is_permutation = True
+            elif data.dtype.kind == "i":
+                is_permutation = data.min(initial=0) >= 0
+            else:
+                n = data.shape[-1]
+                is_permutation = np.count_nonzero(data) == (data.size // n if n else 0)
         result = FactState.TRUE if is_permutation else FactState.FALSE
 
     var.tag.is_permutation = result

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/assumptions/selection.py

@@ -28,12 +28,18 @@ def _selection_from_constant(var: TensorConstant) -> FactState:
     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())
+        with np.errstate(invalid="ignore"):
+            if not (data.sum(axis=-2) == 1).all():
+                is_selection = False
+            elif data.dtype.kind in "ub":
+                is_selection = True
+            elif data.dtype.kind == "i":
+                is_selection = data.min(initial=0) >= 0
+            else:
+                n_rows = data.shape[-2]
+                is_selection = np.count_nonzero(data) == (
+                    data.size // n_rows if n_rows else 0
+                )
         result = FactState.TRUE if is_selection else FactState.FALSE
 
     var.tag.is_selection = result

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/compile/debug/profiling.py

@@ -15,6 +15,7 @@ import operator
 import sys
 import time
 from collections import Counter, defaultdict
+from collections.abc import Callable
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Any
@@ -172,7 +173,7 @@ def print_global_stats():
         print("=" * 50, file=f)
 
 
-_profiler_printers = []
+_profiler_printers: list[Callable] = []
 
 
 def register_profiler_printer(fct):

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/graph/basic.py

@@ -820,40 +820,35 @@ def _make_frozen_output_reduce(out: Variable):
 class FrozenApply(Apply):
     """An immutable, globally-interned Apply node for frozen graphs.
 
-    Uses tuples for ``inputs`` and ``outputs`` so mutation raises ``TypeError``
-    at the language level.  Interned by ``(op, cache_key(inputs))`` —
-    constructing a ``FrozenApply`` with the same op and input variables returns
-    the cached instance.
-
-    Constants are keyed by their ``signature()`` so that two independently
-    created Constants with the same value resolve to the same cached node.
+    ``inputs`` and ``outputs`` are tuples, so mutating them raises ``TypeError``.
+
+    Instances are interned on ``(op, inputs, output_types)``: constructing one
+    with a matching key returns the cached instance.  Constant inputs are keyed
+    by ``signature()`` (so equal-valued Constants share a node).  Other inputs
+    are already globally interned, so identity is enough; the key stores their
+    ``id()`` rather than the variables themselves, keeping strong references out
+    of the cache so chains of ``FrozenApply`` nodes collect in a single GC pass.
+
+    ``output_types`` is in the key because frozen graphs root on
+    ``NominalVariable`` inputs (index and type only).  Nominalizing truncates the
+    root inputs, discarding the values and subgraphs an Op may have read to
+    derive its output type (e.g. via ``infer_shape``), so ``(op, inputs)`` alone
+    no longer pins it down.
     """
 
     _cache: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
 
-    @staticmethod
-    def _input_to_key(inp: Variable):
-        """Convert an input Variable to a hashable, value-based cache key element.
-
-        Non-Constants (NominalVariables, FrozenApply outputs) are already
-        globally interned, so ``id()`` is a correct identity key.  Using
-        ``id()`` instead of the object itself avoids strong references in
-        cache keys that would prevent GC from collecting chains of
-        FrozenApply nodes in a single pass.
-
-        Constants use their ``signature()`` for value-based deduplication.
-        """
-        if isinstance(inp, Constant):
-            return inp.signature()
-        return id(inp)
-
     def __new__(
         cls,
         op: "Op",
         inputs: tuple[Variable, ...],
         output_types: tuple["Type", ...],
     ):
-        cache_key = (op, tuple(cls._input_to_key(i) for i in inputs))
+        cache_key = (
+            op,
+            tuple(i.signature() if isinstance(i, Constant) else id(i) for i in inputs),
+            output_types,
+        )
         cached = cls._cache.get(cache_key)
         if cached is not None:
             return cached

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/graph/fg.py

@@ -1074,12 +1074,17 @@ class FrozenFunctionGraph(AbstractFunctionGraph):
             self._clients = clients
         return self._clients
 
-    def bind(self, replace: dict[Variable, Variable]) -> list[Variable]:
+    def bind(
+        self, replace: Sequence[Variable] | dict[Variable, Variable]
+    ) -> list[Variable]:
         """Return fresh outputs with root inputs substituted per *replace*.
 
         Constants are reused; any non-Constant input not in *replace* raises KeyError.
         """
-        memo = replace.copy()
+        if isinstance(replace, dict):
+            memo = replace.copy()
+        else:
+            memo = dict(zip(self.inputs, replace))
         for node in self.toposort():
             for i in node.inputs:
                 if i not in memo:

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/link/numba/dispatch/elemwise.py

@@ -794,7 +794,7 @@ def numba_funcify_IndexedElemwise(op, node, **kwargs):
 
         return impl
 
-    cache_version = 1
+    cache_version = 2
     if scalar_cache_key is None:
         key = None
     else:

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/link/numba/dispatch/sparse/math.py

@@ -136,7 +136,9 @@ def numba_funcify_SparseDot(op, node, **kwargs):
     x_format = x.type.format if x_is_sparse else None
     y_format = y.type.format if y_is_sparse else None
 
-    cache_version = 2
+    out_type = np.dtype(out_dtype).type
+
+    cache_version = 3
     cache_key = sha256(
         str(
             (
@@ -371,7 +373,7 @@ def numba_funcify_SparseDot(op, node, **kwargs):
         for row_idx in range(n):
             for idx in range(x_ptr[row_idx], x_ptr[row_idx + 1]):
                 col_idx = x_ind[idx]
-                value = x_data[idx]
+                value = out_type(x_data[idx])
                 z[row_idx] += value * y[col_idx]
         return z
 
@@ -390,7 +392,7 @@ def numba_funcify_SparseDot(op, node, **kwargs):
         for col_idx in range(p):
             for idx in range(x_ptr[col_idx], x_ptr[col_idx + 1]):
                 row_idx = x_ind[idx]
-                value = x_data[idx]
+                value = out_type(x_data[idx])
                 z[row_idx] += value * y[col_idx]
         return z
 

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -964,32 +964,30 @@ def _vectorized(
             if spec is None:
                 continue
             indexed_axes = {src_axis: idx_k for idx_k, src_axis in spec}
-            n_indexed = len(indexed_axes)
             n_index_loop_dims = max(idx_types[idx_k].ndim for idx_k, _ in spec)
-            if n_indexed == n_index_loop_dims:
-                iter_shapes[p] = list(iter_shapes[p])
-                for idx_k, axis in spec:
-                    iter_shapes[p][axis] = idx_shapes[idx_k][0]
-            else:
-                batch_ndim = len(input_bc_patterns[p])
-                max_axis = max(a for _, a in spec)
-                new_shape = []
-                new_bc = []
-                src_d = 0
-                idx_d = 0
-                for loop_d in range(batch_ndim):
-                    if src_d in indexed_axes and idx_d < n_index_loop_dims:
-                        new_shape.append(one)
-                        new_bc.append(True)
-                        idx_d += 1
-                        if idx_d >= n_index_loop_dims:
-                            src_d = max_axis + 1
-                    else:
-                        new_shape.append(in_shapes[p][src_d])
-                        new_bc.append(iter_bc[p][loop_d])
-                        src_d += 1
-                iter_shapes[p] = new_shape
-                iter_bc[p] = tuple(new_bc)
+            # An indexed input imposes no constraint on the loop dims produced by
+            # its indices: those are pinned by the index arrays' own iter_shape
+            # entries below.  Mark each such loop dim broadcastable and copy the
+            # source shape for the non-indexed dims.
+            batch_ndim = len(input_bc_patterns[p])
+            max_axis = max(a for _, a in spec)
+            new_shape = []
+            new_bc = []
+            src_d = 0
+            idx_d = 0
+            for loop_d in range(batch_ndim):
+                if src_d in indexed_axes and idx_d < n_index_loop_dims:
+                    new_shape.append(one)
+                    new_bc.append(True)
+                    idx_d += 1
+                    if idx_d >= n_index_loop_dims:
+                        src_d = max_axis + 1
+                else:
+                    new_shape.append(in_shapes[p][src_d])
+                    new_bc.append(iter_bc[p][loop_d])
+                    src_d += 1
+            iter_shapes[p] = new_shape
+            iter_bc[p] = tuple(new_bc)
 
         # Each index array participates in iter_shape validation.
         # Write indices can broadcast against each other, but if ALL write

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/sparse/math.py

@@ -11,8 +11,10 @@ import pytensor.tensor.math as ptm
 from pytensor import config
 from pytensor.gradient import grad_not_implemented
 from pytensor.graph import Apply, Op
+from pytensor.graph.replace import _vectorize_node
 from pytensor.link.c.op import COp
 from pytensor.sparse.type import SparseTensorType
+from pytensor.tensor.reshape import join_dims, split_dims
 from pytensor.tensor.shape import specify_broadcastable
 from pytensor.tensor.type import TensorType, Variable, complex_dtypes, tensor
 
@@ -2067,3 +2069,49 @@ class Usmm(Op):
 
 
 usmm = Usmm()
+
+
+@_vectorize_node.register(StructuredDot)
+def _vectorize_structured_dot(op, node, batch_a, batch_b):
+    """Batch StructuredDot(sparse_const, dense): (m,k)@(B...,k,n) -> (B...,m,n).
+
+    The sparse left input must stay unbatched (scipy has no batched-sparse
+    type). The dense right input may gain any number of leading batch dims;
+    we fold them through the existing 2D StructuredDot with a moveaxis +
+    join_dims / split_dims round trip.
+    """
+    a, b = node.inputs
+    if batch_a.type.ndim != a.type.ndim:
+        raise NotImplementedError(
+            "Cannot vectorize StructuredDot when the sparse (left) input is "
+            "batched; scipy has no batched-sparse type."
+        )
+
+    if batch_b.type.ndim == b.type.ndim:
+        # No batch dims added to the dense input — rebuild the op as-is.
+        return op.make_node(batch_a, batch_b).outputs
+
+    # batch_b is (B1,...,BN, k, n). Move k to the front and fold the batch
+    # dims and n into a single column axis: (k, B1*...*BN*n).
+    moved = ptb.moveaxis(batch_b, -2, 0)  # (k, B1,...,BN, n)
+    trailing = moved.shape[1:]
+    flat_b = join_dims(moved, start_axis=1)  # (k, B*n)
+
+    flat_out = op.make_node(batch_a, flat_b).outputs[0]  # (m, B*n)
+
+    # Unfold the column axis and move m back into the (-2) slot.
+    unflat = split_dims(flat_out, shape=trailing, axis=1)  # (m, B1,...,BN, n)
+    out = ptb.moveaxis(unflat, 0, -2)  # (B1,...,BN, m, n)
+    return [out]
+
+
+def _vectorize_sparse_unsupported(op, node, *batched_inputs):
+    raise NotImplementedError(
+        f"Cannot vectorize {type(op).__name__}: scipy has no batched-sparse "
+        "representation, so the sparse operand cannot be broadcast against a "
+        "batched dense input."
+    )
+
+
+for _op_cls in (AddSD, SparseDenseMultiply):
+    _vectorize_node.register(_op_cls)(_vectorize_sparse_unsupported)

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/tensor/rewriting/linalg/summary.py

@@ -18,6 +18,7 @@ from pytensor.tensor.linalg.decomposition.cholesky import Cholesky
 from pytensor.tensor.linalg.decomposition.lu import LU, LUFactor
 from pytensor.tensor.linalg.decomposition.qr import QR
 from pytensor.tensor.linalg.decomposition.svd import SVD
+from pytensor.tensor.linalg.inverse import MatrixInverse
 from pytensor.tensor.linalg.summary import SLogDet, det
 from pytensor.tensor.math import Prod, log, prod
 from pytensor.tensor.rewriting.basic import (
@@ -192,6 +193,16 @@ def det_of_triangular(fgraph, node):
     return [det_val]
 
 
+@register_canonicalize
+@register_stabilize
+@node_rewriter([det])
+def det_of_inv(fgraph, node):
+    """Replace det(matrix_inverse(X)) with reciprocal(det(X))."""
+    match node.inputs[0].owner_op_and_inputs:
+        case (Blockwise(MatrixInverse()), X):
+            return [1 / det(X)]
+
+
 @register_specialize
 @node_rewriter([det])
 def slogdet_specialization(fgraph, node):

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/tensor/rewriting/math.py

@@ -112,7 +112,7 @@ from pytensor.tensor.rewriting.basic import (
 from pytensor.tensor.rewriting.blockwise import blockwise_of
 from pytensor.tensor.rewriting.elemwise import apply_local_dimshuffle_lift
 from pytensor.tensor.shape import Shape, Shape_i, specify_shape
-from pytensor.tensor.subtensor import Subtensor
+from pytensor.tensor.subtensor import Subtensor, _is_provably_positive
 from pytensor.tensor.type import (
     complex_dtypes,
     uint_dtypes,
@@ -689,6 +689,66 @@ def local_exp_log_nan_switch(fgraph, node):
         return [new_out]
 
 
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([log])
+def local_log_div(fgraph, node):
+    """Rewrite log(reciprocal(x)) -> -log(x) and log(a / b) -> log(a) - log(b).
+
+    A reciprocal is just ``1 / x``; log(a / b) only splits when a positive
+    constant operand is involved, so its log folds and the op count stays flat.
+    """
+    (inp,) = node.inputs
+    if not (inp.owner and isinstance(inp.owner.op, Elemwise)):
+        return None
+    scalar_op = inp.owner.op.scalar_op
+
+    if isinstance(scalar_op, ps.Reciprocal):
+        return [neg(log(inp.owner.inputs[0]))]
+
+    if isinstance(scalar_op, ps.TrueDiv):
+        num, den = inp.owner.inputs
+        if (isinstance(num, Constant) and _is_provably_positive(num, strict=True)) or (
+            isinstance(den, Constant) and _is_provably_positive(den, strict=True)
+        ):
+            return [log(num) - log(den)]
+
+
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([sign])
+def local_sign_div(fgraph, node):
+    """Rewrite sign of a reciprocal or division from a known-sign operand.
+
+    sign(reciprocal(x)) -> sign(x). For sign(a / b): a provably positive side ->
+    ``sign(other)``; a negative constant side -> ``-sign(other)``. Bails out
+    otherwise.
+    """
+    (inp,) = node.inputs
+    if not (inp.owner and isinstance(inp.owner.op, Elemwise)):
+        return None
+    scalar_op = inp.owner.op.scalar_op
+
+    if isinstance(scalar_op, ps.Reciprocal):
+        return [sign(inp.owner.inputs[0])]
+
+    if not isinstance(scalar_op, ps.TrueDiv):
+        return None
+
+    num, den = inp.owner.inputs
+
+    if _is_provably_positive(num, strict=True):
+        return [sign(den)]
+    if _is_provably_positive(den, strict=True):
+        return [sign(num)]
+
+    for side, other in ((num, den), (den, num)):
+        if isinstance(side, Constant) and np.all(np.asarray(side.data) < 0):
+            return [neg(sign(other))]
+
+
 @register_canonicalize
 @register_specialize
 @node_rewriter([Sum])

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/tensor/rewriting/subtensor.py

@@ -2738,22 +2738,26 @@ def local_blockwise_inc_subtensor(fgraph, node):
     idxs_core_ndim = [len(inp_sig) for inp_sig in node.op.inputs_sig[2:]]
     max_idx_core_ndim = max(idxs_core_ndim, default=0)
 
-    # Broadcast buffer to batch_shape
-    if x.type.broadcastable != out.type.broadcastable:
-        batch_shape = [1] * batch_ndim
-        for inp in node.inputs:
-            for i, (broadcastable, batch_dim) in enumerate(
-                zip(inp.type.broadcastable[:batch_ndim], tuple(inp.shape)[:batch_ndim])
-            ):
-                if broadcastable:
-                    # This dimension is broadcastable, it doesn't provide shape information
-                    continue
-                if batch_shape[i] != 1:
-                    # We already found a source of shape for this batch dimension
-                    continue
-                batch_shape[i] = batch_dim
+    # Broadcast buffer to batch_shape. The output batch shape and broadcast
+    # pattern are derived from the inputs, never from `out.type`, which can be
+    # stale after an upstream rewrite swaps an input.
+    batch_shape = [1] * batch_ndim
+    out_batch_bcast = [True] * batch_ndim
+    for inp in node.inputs:
+        for i, (broadcastable, batch_dim) in enumerate(
+            zip(inp.type.broadcastable[:batch_ndim], tuple(inp.shape)[:batch_ndim])
+        ):
+            if broadcastable:
+                # This dimension is broadcastable, it doesn't provide shape information
+                continue
+            out_batch_bcast[i] = False
+            if batch_shape[i] != 1:
+                # We already found a source of shape for this batch dimension
+                continue
+            batch_shape[i] = batch_dim
+
+    if list(x.type.broadcastable[:batch_ndim]) != out_batch_bcast:
         x = broadcast_to(x, (*batch_shape, *x.shape[batch_ndim:]))
-        assert x.type.broadcastable == out.type.broadcastable
 
     # Massage indices so they respect blockwise semantics while using regular indexing
     core_idxs = []

{pytensor-3.0.4 → pytensor-3.0.5}/pytensor/tensor/rewriting/subtensor_lift.py

@@ -567,8 +567,13 @@ def local_subtensor_of_expand_dims(fgraph, node):
 
     idx_tuple = indices_from_subtensor(idx, node.op.idx_list)
 
-    # Keep indexes for the original dimensions, and drop indexes for the expanded dimensions when safe
+    # Keep indexes for the original dimensions, and drop indexes for the expanded
+    # dimensions when safe. We also track where each kept expanded dimension lands
+    # in the output (`out_pos`), so it can be re-introduced with expand_dims below.
+    # These axes are derived from the indices, not the (possibly stale) output type.
     new_idxs = []
+    expand_axes = []
+    out_pos = 0
     for i, idx_item in enumerate(idx_tuple):
         if i in expanded_axes:
             if isinstance(idx_item, slice):
@@ -576,6 +581,8 @@ def local_subtensor_of_expand_dims(fgraph, node):
                 if idx_item == slice(None):
                     # A None slice, always keeps the dimension.
                     # We skip the index, and later introduce the needed expand_dim
+                    expand_axes.append(out_pos)
+                    out_pos += 1
                     continue
                 else:
                     # Other slices could keep or drop the dimension.
@@ -589,27 +596,23 @@ def local_subtensor_of_expand_dims(fgraph, node):
         else:
             # Keep indexes for non-expanded dimensions
             new_idxs.append(idx_item)
+            # An integer index drops the dimension; any slice keeps one.
+            if isinstance(idx_item, slice):
+                out_pos += 1
+
+    # Trailing dimensions beyond the explicit indices are implicit full slices;
+    # the expanded ones among them must also be re-introduced.
+    for axis in range(len(idx_tuple), ds.type.ndim):
+        if axis in expanded_axes:
+            expand_axes.append(out_pos)
+        out_pos += 1
 
     [old_out] = node.outputs
     out = x[tuple(new_idxs)]
+    if expand_axes:
+        out = expand_dims(out, axis=expand_axes)
     copy_stack_trace(old_out, out)
 
-    if out.type.broadcastable != old_out.type.broadcastable:
-        # Re-introduce needed new dimensions (corresponding to full slices on the original expanded dimensions)
-        # If out.type.broadcastable == (False) and old_out.type.broadcastable == (True, False, True)
-        # then axis = (0, 2)
-        old_bcast = list(old_out.type.broadcastable)
-        expanded_bcast = list(out.type.broadcastable)
-        axis = []
-        i = 0
-        while i < len(old_bcast):
-            if i == len(expanded_bcast) or expanded_bcast[i] != old_bcast[i]:
-                expanded_bcast.insert(i, True)
-                axis.append(i)
-            i += 1
-        out = expand_dims(out, axis=axis)
-        copy_stack_trace(old_out, out)
-
     return [out]
 
 

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

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