pytensor 3.0.1__tar.gz → 3.1.0__tar.gz

{pytensor-3.0.1/pytensor.egg-info → pytensor-3.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.0.1
+Version: 3.1.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

{pytensor-3.0.1 → pytensor-3.1.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-05-07T19:28:03+0200",
+ "date": "2026-05-20T17:31:39-0500",
  "dirty": false,
  "error": null,
- "full-revisionid": "10852419d00cf04f0c89c034f04ddfa9e5067e63",
- "version": "3.0.1"
+ "full-revisionid": "b712ea6e6a7541ce24c47f8869699178958dadf6",
+ "version": "3.1.0"
 }
 '''  # END VERSION_JSON
 

pytensor-3.1.0/pytensor/assumptions/__init__.py

@@ -0,0 +1,88 @@
+# Import every rule module so its registrations land on the global registry that
+# ``check_assumption`` consults. The property modules pull in the shared helper
+# modules (alloc, dimshuffle, subtensor, dot, elemwise) transitively.
+import pytensor.assumptions.alloc
+import pytensor.assumptions.blockwise
+import pytensor.assumptions.diagonal
+import pytensor.assumptions.dimshuffle
+import pytensor.assumptions.orthogonal
+import pytensor.assumptions.positive_definite
+import pytensor.assumptions.reshape
+import pytensor.assumptions.shape
+import pytensor.assumptions.subtensor
+import pytensor.assumptions.symmetric
+import pytensor.assumptions.triangular
+from pytensor.assumptions.core import (
+    ALL_KEYS,
+    DIAGONAL,
+    IMPLIES,
+    LOWER_TRIANGULAR,
+    ORTHOGONAL,
+    POSITIVE_DEFINITE,
+    SYMMETRIC,
+    UPPER_TRIANGULAR,
+    AssumptionFeature,
+    AssumptionKey,
+    ConflictingAssumptionsError,
+    FactState,
+    check_assumption,
+    register_assumption,
+    register_constant_inference,
+    register_implies,
+)
+from pytensor.assumptions.specify import (
+    SpecifyAssumptions,
+    assume,
+    specify_assumption_rule,
+)
+from pytensor.graph.fg import FunctionGraph
+
+
+def summarize_assumptions(feature: AssumptionFeature, var) -> str:
+    """Return a compact ``{...}`` tag of the facts known about *var*.
+
+    Each TRUE fact appears as its ``short_name``; each FALSE fact as ``!short_name``.
+    Implied facts are pruned: a TRUE fact is dropped when a stronger fact (one that
+    implies it) is also TRUE, and a FALSE fact is dropped when a weaker fact (one it
+    implies) is also FALSE. Return ``""`` when nothing is known.
+    """
+    states = {}
+    for key in ALL_KEYS:
+        try:
+            states[key] = feature.get(var, key)
+        except ConflictingAssumptionsError:
+            states[key] = FactState.CONFLICT
+
+    tokens = []
+    for key in ALL_KEYS:
+        state = states[key]
+        label = key.short_name or key.name
+        if state is FactState.TRUE:
+            if any(
+                key in weaker and states[stronger] is FactState.TRUE
+                for stronger, weaker in IMPLIES.items()
+            ):
+                continue
+            tokens.append(label)
+        elif state is FactState.FALSE:
+            if any(states[w] is FactState.FALSE for w in IMPLIES.get(key, ())):
+                continue
+            tokens.append(f"!{label}")
+        elif state is FactState.CONFLICT:
+            tokens.append(f"{label}=CONFLICT")
+
+    return "{" + ", ".join(tokens) + "}" if tokens else ""
+
+
+def assumption_tags(outputs) -> dict:
+    """Map each variable reachable from *outputs* to its :func:`summarize_assumptions`
+    tag, omitting variables with no known facts."""
+    fgraph = FunctionGraph(outputs=list(outputs), clone=False)
+    feature = AssumptionFeature()
+    fgraph.attach_feature(feature)
+    tags = {}
+    for var in fgraph.variables:
+        tag = summarize_assumptions(feature, var)
+        if tag:
+            tags[var] = tag
+    return tags

pytensor-3.1.0/pytensor/assumptions/alloc.py

@@ -0,0 +1,90 @@
+from pytensor.assumptions.core import (
+    ALL_KEYS,
+    FactState,
+    register_assumption,
+    true_if,
+)
+from pytensor.tensor.basic import (
+    Alloc,
+    NotScalarConstantError,
+    get_underlying_scalar_constant_value,
+)
+from pytensor.tensor.variable import TensorConstant
+
+
+def alloc_of_zero(key, op, feature, fgraph, node, input_states) -> list[FactState]:
+    """``Alloc`` rule for DIAGONAL / LOWER_TRIANGULAR / UPPER_TRIANGULAR: TRUE when
+    the fill value is the scalar 0 (an all-zero square matrix), FALSE when it is a
+    known non-zero scalar -- the off-diagonal entries are then non-zero, so none of
+    these properties holds.
+
+    Requires the trailing two output dims to be statically known and equal; these
+    properties apply only to square matrices. SYMMETRIC uses :func:`alloc_is_symmetric`
+    instead, since a constant fill is symmetric whatever its value.
+    """
+    out_shape = node.outputs[0].type.shape
+    if len(out_shape) < 2:
+        return [FactState.UNKNOWN]
+    m, n = out_shape[-2], out_shape[-1]
+    if m is None or n is None or m != n:
+        return [FactState.UNKNOWN]
+    try:
+        val = get_underlying_scalar_constant_value(node.inputs[0])
+    except NotScalarConstantError:
+        return [FactState.UNKNOWN]
+    return true_if(val == 0, else_false=True)
+
+
+def alloc_is_symmetric(key, op, feature, fgraph, node, input_states) -> list[FactState]:
+    """``Alloc`` rule for SYMMETRIC: a scalar fill broadcasts one value to every
+    entry, so a square output equals its transpose -- whatever the value.
+
+    A non-scalar fill (e.g. a vector broadcast across rows) is left UNKNOWN.
+    """
+    out_shape = node.outputs[0].type.shape
+    if len(out_shape) < 2:
+        return [FactState.UNKNOWN]
+    m, n = out_shape[-2], out_shape[-1]
+    if m is None or n is None or m != n:
+        return [FactState.UNKNOWN]
+    return true_if(node.inputs[0].type.ndim == 0)
+
+
+def eye_identity_rule(key, op, feature, fgraph, node, input_states) -> list[FactState]:
+    """Rule body: TRUE when an :class:`Eye` node produces the identity matrix (square, k == 0)."""
+    n, m, k = node.inputs
+    if not isinstance(k, TensorConstant):
+        return [FactState.UNKNOWN]
+    if k.data.item() != 0:
+        # Ones sit on an off-main diagonal, so this is not the identity.
+        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)
+    return [FactState.UNKNOWN]
+
+
+def alloc_diag_at_offset_zero(
+    key, op, feature, fgraph, node, input_states
+) -> list[FactState]:
+    """Rule body: TRUE when :class:`AllocDiag` places values on the main diagonal,
+    FALSE when it places them on any other diagonal (off-main entries break diagonal /
+    symmetric / PD structure regardless of the diagonal vector's values)."""
+    return [FactState.TRUE if op.offset == 0 else FactState.FALSE]
+
+
+def alloc_propagates_matrix_property(
+    key, op, feature, fgraph, node, input_states
+) -> list[FactState]:
+    """``Alloc`` broadcasts a value across new or expanded leading axes. When the
+    value is already matrix-shaped with non-broadcastable core axes, the alloc
+    cannot touch those axes, so the matrix property carries through unchanged."""
+    value = node.inputs[0]
+    if value.type.ndim >= 2 and not any(value.type.broadcastable[-2:]):
+        return [input_states[0]]
+    return [FactState.UNKNOWN]
+
+
+for _key in ALL_KEYS:
+    register_assumption(_key, Alloc)(alloc_propagates_matrix_property)

pytensor-3.1.0/pytensor/assumptions/blockwise.py

@@ -0,0 +1,17 @@
+from pytensor.assumptions.core import (
+    ALL_KEYS,
+    infer_assumption_for_node,
+    register_assumption,
+)
+from pytensor.tensor.blockwise import Blockwise
+
+
+def _blockwise_delegate(key, op, feature, fgraph, node, input_states):
+    """Delegate assumption inference to the ``core_op`` of a Blockwise wrapper."""
+    return infer_assumption_for_node(
+        key, op.core_op, feature, fgraph, node, input_states
+    )
+
+
+for _key in ALL_KEYS:
+    register_assumption(_key, Blockwise)(_blockwise_delegate)

pytensor-3.1.0/pytensor/assumptions/core.py

@@ -0,0 +1,373 @@
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import IntFlag, auto
+from typing import Any
+
+from pytensor.graph import Apply, FunctionGraph, Op
+from pytensor.graph.features import AlreadyThere, Feature
+from pytensor.graph.traversal import walk_toposort
+from pytensor.tensor.variable import TensorConstant
+
+
+class FactState(IntFlag):
+    """Three-valued logic for assumption inference.
+
+    The three fact states are TRUE, FALSE and UNKNOWN.
+
+    UNKNOWN is the default condition, in which we cannot confirm or deny the fact.  TRUE and FALSE are definitive
+    states.  If we have evidence that a fact is both TRUE and FALSE, we get CONFLICT. In general, a CONFLICT state
+    should not be possible.
+    """
+
+    UNKNOWN = 0
+    TRUE = auto()
+    FALSE = auto()
+    CONFLICT = TRUE | FALSE
+
+    def __bool__(self) -> bool:
+        raise TypeError(
+            "FactState has no boolean value; compare against FactState.TRUE, "
+            "FactState.UNKNOWN, etc. explicitly, or use AssumptionFeature.check()."
+        )
+
+    @classmethod
+    def join(cls, left: "FactState", right: "FactState") -> "FactState":
+        """Combine two pieces of evidence about the *same* (variable, key)."""
+        return cls(left | right)
+
+
+@dataclass(frozen=True)
+class AssumptionKey:
+    """Identifies a named structural property (e.g. "diagonal" or "triangular").
+
+    ``short_name`` is an abbreviated label used by ``debugprint(print_assumptions=True)``;
+    it falls back to ``name`` when empty.
+    """
+
+    name: str
+    short_name: str = ""
+
+    def __repr__(self) -> str:
+        return self.name
+
+
+class ConflictingAssumptionsError(ValueError):
+    """Raised when joining evidence about a (variable, key) produces ``FactState.CONFLICT``.
+
+    Surfaces contradictions between user-asserted facts and inferred facts (or between
+    multiple sources of inferred facts) before they can silently corrupt downstream rewrites.
+    """
+
+
+# An inference function takes the active AssumptionKey, the Op, the current AssumptionFeature, the FunctionGraph,
+# the Apply node being analyzed, and the states of the input variables for the current key. It returns a list of
+# FactState with one entry per node output.
+InferFactFn = Callable[
+    [AssumptionKey, Op, "AssumptionFeature", FunctionGraph, Apply, list[FactState]],
+    list[FactState],
+]
+
+# The global inference registry maps (AssumptionKey, Op type) pairs to lists of inference functions.
+# Rules are tried in registration order; the first to return TRUE wins.
+ASSUMPTION_INFER_REGISTRY: dict[tuple[AssumptionKey, type], list[InferFactFn]] = {}
+
+# Registry mapping assumptions to other assumptions they imply.  For example, a "diagonal" matrix is also "symmetric"
+# and "triangular".  This is consulted after all other inference rules to derive additional facts.
+IMPLIES: dict[AssumptionKey, list[AssumptionKey]] = {}
+
+
+def register_implies(stronger: AssumptionKey, *weaker: AssumptionKey) -> None:
+    """Declare that *stronger* being TRUE implies each *weaker* key is also TRUE."""
+    IMPLIES.setdefault(stronger, []).extend(weaker)
+
+
+# Maps a key to a function that infers it directly from a literal TensorConstant's data.
+# AssumptionFeature consults this for variables with no owner, letting a property be
+# derived from raw values (e.g. recognising a constant diagonal matrix).
+ConstantInferFn = Callable[[TensorConstant], FactState]
+CONSTANT_INFER_REGISTRY: dict[AssumptionKey, ConstantInferFn] = {}
+
+
+def register_constant_inference(key: AssumptionKey, fn: ConstantInferFn) -> None:
+    """Register *fn* as the literal-constant inference for *key*."""
+    CONSTANT_INFER_REGISTRY[key] = fn
+
+
+# The canonical structural-property keys
+DIAGONAL = AssumptionKey("diagonal", short_name="diag")
+LOWER_TRIANGULAR = AssumptionKey("lower_triangular", short_name="tril")
+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")
+
+ALL_KEYS = (
+    DIAGONAL,
+    LOWER_TRIANGULAR,
+    UPPER_TRIANGULAR,
+    SYMMETRIC,
+    POSITIVE_DEFINITE,
+    ORTHOGONAL,
+)
+
+# Implications about structural properties derivably from other structural properties
+register_implies(DIAGONAL, LOWER_TRIANGULAR, UPPER_TRIANGULAR, SYMMETRIC)
+register_implies(POSITIVE_DEFINITE, SYMMETRIC)
+
+
+def register_assumption(
+    key: AssumptionKey, *op_types: type
+) -> Callable[[InferFactFn], InferFactFn]:
+    """Decorator that registers an inference rule for ``(key, op_type)`` pairs.
+
+    The decorated function is called as ``fn(key, op, feature, fgraph, node, input_states)``
+    and must return a list of :class:`FactState` with one entry per node output.
+    """
+
+    def decorator(fn: InferFactFn) -> InferFactFn:
+        for op_type in op_types:
+            ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), []).append(fn)
+        return fn
+
+    return decorator
+
+
+def infer_assumption_for_node(
+    key: AssumptionKey,
+    op: Op,
+    feature: "AssumptionFeature",
+    fgraph: FunctionGraph,
+    node: Apply,
+    input_states: list[FactState],
+) -> list[FactState]:
+    """Determine the *key* fact for every output of *node*.
+
+    Walks the Op's MRO for the most specific :func:`register_assumption` rule and
+    returns the first non-UNKNOWN result. Falls back to ``UNKNOWN`` per output.
+    """
+    rules: list[InferFactFn] = []
+    for cls in type(op).__mro__:
+        fns = ASSUMPTION_INFER_REGISTRY.get((key, cls))
+        if fns is not None:
+            rules = fns
+            break
+
+    n_outputs = len(node.outputs)
+    for fn in rules:
+        output_states = fn(key, op, feature, fgraph, node, input_states)
+        if len(output_states) != n_outputs:
+            raise ValueError(
+                f"an assumption rule returned {len(output_states)} states for "
+                f"{n_outputs} outputs on node {node!r}"
+            )
+        if any(s is not FactState.UNKNOWN for s in output_states):
+            return output_states
+
+    return [FactState.UNKNOWN] * n_outputs
+
+
+class AssumptionFeature(Feature):
+    """``FunctionGraph`` feature that tracks symbolic assumptions about variables.
+
+    Assumptions (e.g. "this matrix is diagonal") are represented as ``(variable,
+    AssumptionKey) -> FactState`` mappings. Facts are inferred lazily via per-Op
+    rules registered with :func:`register_assumption`.
+
+    Attached lazily by :func:`check_assumption` on first use.
+    Results are cached and automatically refreshed when the graph changes.
+    """
+
+    __slots__ = ("_stale_vars", "_var_to_keys", "cache", "fgraph")
+
+    def on_attach(self, fgraph: Any) -> None:
+        if hasattr(fgraph, "assumption_feature"):
+            raise AlreadyThere("AssumptionFeature is already attached")
+        self.fgraph = fgraph
+        self.cache: dict[tuple[Any, AssumptionKey], FactState] = {}
+        self._var_to_keys: dict[Any, set[AssumptionKey]] = {}
+        self._stale_vars: set = set()
+        fgraph.assumption_feature = self
+
+    def on_detach(self, fgraph: Any) -> None:
+        self.cache = {}
+        self._var_to_keys = {}
+        self._stale_vars = set()
+        self.fgraph = None
+        del fgraph.assumption_feature
+
+    def on_change_input(self, fgraph, node, i, old_var, new_var, reason=None) -> None:
+        # Carry non-UNKNOWN facts from old_var to new_var, then defer downstream
+        # UNKNOWN invalidation to the next get()/check() call.
+        # CONFLICT is cached, not raised: the graph mutation has already landed,
+        # so raising here corrupts the caller's view.
+        for key in self._var_to_keys.get(old_var, ()):
+            old_state = self.cache.get((old_var, key))
+            if old_state is None or old_state is FactState.UNKNOWN:
+                continue
+            new_state = self.cache.get((new_var, key))
+            if new_state is None or new_state is FactState.UNKNOWN:
+                self.cache[(new_var, key)] = old_state
+            else:
+                self.cache[(new_var, key)] = FactState.join(old_state, new_state)
+            self._var_to_keys.setdefault(new_var, set()).add(key)
+
+        self._stale_vars.update(node.outputs)
+
+    def on_prune(self, fgraph, node, reason=None) -> None:
+        self._stale_vars.difference_update(node.outputs)
+
+    def clone(self) -> "AssumptionFeature":
+        # Cache entries are keyed on variable identity, so the cloned feature starts empty
+        # and re-derives facts from the cloned graph (assumptions attached via ``assume()``
+        # follow their variables through the clone since they live as graph nodes).
+        return AssumptionFeature()
+
+    def get(self, var: Any, key: AssumptionKey) -> FactState:
+        """Return the inferred :class:`FactState` for ``(var, key)``.
+
+        Raise :class:`ConflictingAssumptionsError` when the cached state is
+        :attr:`FactState.CONFLICT` — surfaces conflicts from substitutions or
+        owner-inferred rules at the point of query, where the graph is stable
+        and the caller can react.
+        """
+        cache_key = (var, key)
+        state = self.cache.get(cache_key)
+        if state is None or state is FactState.UNKNOWN:
+            if self._stale_vars:
+                self._drop_unknown_downstream(self._stale_vars)
+                self._stale_vars.clear()
+                # Re-check: the UNKNOWN will still be valid
+                # if the stale vars were not ancestors of var.
+                state = self.cache.get(cache_key)
+            if state is None:
+                state = self._compute(var, key)
+        if state is FactState.CONFLICT:
+            raise ConflictingAssumptionsError(
+                f"Conflicting evidence for key {key!r} on {var!r}."
+            )
+        return state
+
+    def check(self, var: Any, key: AssumptionKey) -> bool:
+        """Return ``True`` iff the assumption is definitively TRUE for ``var``."""
+        return self.get(var, key) is FactState.TRUE
+
+    def _compute(self, var: Any, key: AssumptionKey) -> FactState:
+        """Infer the fact state for ``(var, key)``, caching ancestors inputs-first.
+
+        After owner-based inference, walks the ``IMPLIES`` graph in both directions:
+          - Forward: a stronger key being TRUE makes this (weaker) key TRUE.
+          - Contrapositive: a weaker key being FALSE makes this (stronger) key FALSE.
+
+        Raise :class:`ConflictingAssumptionsError` if owner-inferred rules produce
+        ``FactState.CONFLICT``.
+        """
+
+        def deps(v: Any) -> tuple:
+            # Stop descending at cached variables — their facts (and their ancestors')
+            # are already settled for this key.
+            if (v, key) in self.cache:
+                return ()
+            owner = getattr(v, "owner", None)
+            return tuple(owner.inputs) if owner is not None else ()
+
+        for v in walk_toposort([var], deps):
+            if (v, key) in self.cache:
+                continue
+
+            # Cache UNKNOWN up front so any recursive ``self.get(v, other_key)`` triggered
+            # by the IMPLIES checks sees a placeholder instead of recursing back here.
+            self.cache[(v, key)] = FactState.UNKNOWN
+            self._var_to_keys.setdefault(v, set()).add(key)
+
+            owner = getattr(v, "owner", None)
+            if owner is not None:
+                input_states = [self.get(inp, key) for inp in owner.inputs]
+                output_states = infer_assumption_for_node(
+                    key, owner.op, self, self.fgraph, owner, input_states
+                )
+                state = output_states[owner.outputs.index(v)]
+            elif isinstance(v, TensorConstant) and key in CONSTANT_INFER_REGISTRY:
+                state = CONSTANT_INFER_REGISTRY[key](v)
+            else:
+                state = FactState.UNKNOWN
+
+            if state is FactState.UNKNOWN:
+                for stronger, weaker_list in IMPLIES.items():
+                    if key in weaker_list and self.get(v, stronger) is FactState.TRUE:
+                        state = FactState.TRUE
+                        break
+
+            if state is FactState.UNKNOWN:
+                for weaker in IMPLIES.get(key, ()):
+                    if self.get(v, weaker) is FactState.FALSE:
+                        state = FactState.FALSE
+                        break
+
+            if state is FactState.CONFLICT:
+                raise ConflictingAssumptionsError(
+                    f"Conflicting evidence for {key} on {v!r} from owner-inferred rules."
+                )
+
+            self.cache[(v, key)] = state
+
+        return self.cache[(var, key)]
+
+    def _drop_unknown_downstream(self, start_vars) -> None:
+        """Drop UNKNOWN entries for *start_vars* and downstream so they re-probe."""
+        cache = self.cache
+        var_to_keys = self._var_to_keys
+        clients = self.fgraph.clients
+        queue = deque(start_vars)
+        seen = set(start_vars)
+        while queue:
+            var = queue.popleft()
+            keys = var_to_keys.get(var)
+            if keys:
+                for key in [
+                    k for k in keys if cache.get((var, k)) is FactState.UNKNOWN
+                ]:
+                    cache.pop((var, key), None)
+                    keys.discard(key)
+                if not keys:
+                    var_to_keys.pop(var, None)
+            for client_node, _ in clients.get(var, ()):
+                for out in client_node.outputs:
+                    if out not in seen:
+                        seen.add(out)
+                        queue.append(out)
+
+
+def true_if(cond: bool, else_false: bool = False) -> list[FactState]:
+    """``[TRUE]`` when *cond* holds, ``[UNKNOWN]`` (or ``[FALSE]``) otherwise."""
+    if cond:
+        return [FactState.TRUE]
+    return [FactState.FALSE] if else_false else [FactState.UNKNOWN]
+
+
+def propagate_first(key, op, feature, fgraph, node, input_states) -> list[FactState]:
+    """Output inherits the assumption iff the first input has it."""
+    return true_if(input_states[0] is FactState.TRUE)
+
+
+def all_inputs_have_key(
+    key, op, feature, fgraph, node, input_states
+) -> list[FactState]:
+    """Output inherits the assumption iff *every* input has it."""
+    return true_if(all(s is FactState.TRUE for s in input_states))
+
+
+def check_assumption(
+    fgraph: FunctionGraph | None, var: Any, key: AssumptionKey
+) -> bool:
+    """Return True iff *key* is definitively TRUE for *var* in *fgraph*.
+
+    Lazily attaches :class:`AssumptionFeature` to *fgraph* if it is not already present.
+    """
+    if fgraph is None:
+        return False
+    feature = getattr(fgraph, "assumption_feature", None)
+    if feature is None:
+        feature = AssumptionFeature()
+        fgraph.attach_feature(feature)
+    return feature.check(var, key)