From 789ef24c49f17981d47b0caca7426858df998542 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 22 Aug 2026 22:44:49 -0400 Subject: [PATCH 1/9] Register assumption keys globally and add key-agnostic rules Rules declared with register_universal_assumption install onto keys created after the declaration, so a key defined by a downstream library is no longer silently skipped by the modules that previously looped over a frozen ALL_KEYS tuple at import time. --- pytensor/assumptions/__init__.py | 2 + pytensor/assumptions/blockwise.py | 8 +- pytensor/assumptions/core.py | 105 ++++++++++++++++++++++- pytensor/assumptions/specify.py | 25 +++--- pytensor/tensor/rewriting/assumptions.py | 20 ++--- tests/assumptions/conftest.py | 5 ++ 6 files changed, 130 insertions(+), 35 deletions(-) diff --git a/pytensor/assumptions/__init__.py b/pytensor/assumptions/__init__.py index 8b580139a8..8bfe7c37aa 100644 --- a/pytensor/assumptions/__init__.py +++ b/pytensor/assumptions/__init__.py @@ -18,6 +18,7 @@ ALL_KEYS, DIAGONAL, IMPLIES, + KEY_REGISTRY, LOWER_TRIANGULAR, MATRIX_KEYS, ORTHOGONAL, @@ -35,6 +36,7 @@ register_assumption, register_constant_inference, register_implies, + register_universal_assumption, ) from pytensor.assumptions.specify import ( SpecifyAssumptions, diff --git a/pytensor/assumptions/blockwise.py b/pytensor/assumptions/blockwise.py index 460cf8e3eb..69cba2cdc0 100644 --- a/pytensor/assumptions/blockwise.py +++ b/pytensor/assumptions/blockwise.py @@ -1,17 +1,13 @@ from pytensor.assumptions.core import ( - ALL_KEYS, infer_assumption_for_node, - register_assumption, + register_universal_assumption, ) from pytensor.tensor.blockwise import Blockwise +@register_universal_assumption(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) diff --git a/pytensor/assumptions/core.py b/pytensor/assumptions/core.py index 9b1195fc5d..39ce5d9899 100644 --- a/pytensor/assumptions/core.py +++ b/pytensor/assumptions/core.py @@ -1,5 +1,5 @@ from collections import deque -from collections.abc import Callable +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass from enum import IntFlag, auto from typing import Any @@ -41,13 +41,42 @@ def join(cls, left: "FactState", right: "FactState") -> "FactState": 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. + Constructing a key registers it in :data:`KEY_REGISTRY` and installs every rule + declared with :func:`register_universal_assumption` for it, which is all a + downstream library must do to add a property of its own. + + ``name`` is the key's identity: two keys may not share one unless they are + identical in every field. + + Parameters + ---------- + name : str + Unique identifier for the property. + short_name : str, optional + Abbreviated label used by ``debugprint(print_assumptions=True)``. Falls back + to ``name`` when empty. """ name: str short_name: str = "" + def __post_init__(self) -> None: + registered = KEY_REGISTRY.get(self.name) + if registered is not None: + if registered != self: + raise ValueError( + f"An assumption named {self.name!r} is already registered as " + f"{registered!r} with different metadata. Assumption names are " + f"global identifiers; pick a distinct one." + ) + # Re-creating an identical key (a module imported twice) is a no-op: + # equal keys hash alike, so the installed rules already apply to it. + return + + KEY_REGISTRY[self.name] = self + for op_types, fn in UNIVERSAL_RULES: + _install_rule(self, op_types, fn) + def __repr__(self) -> str: return self.name @@ -72,6 +101,71 @@ class ConflictingAssumptionsError(ValueError): # Rules are tried in registration order; the first to return TRUE wins. ASSUMPTION_INFER_REGISTRY: dict[tuple[AssumptionKey, type], list[InferFactFn]] = {} +# Every AssumptionKey ever constructed, by name. Downstream libraries join the system +# by constructing a key; nothing else is required of them. +KEY_REGISTRY: dict[str, AssumptionKey] = {} + +# Rules that hold for every key regardless of what the key means, as (op_types, fn) +# pairs. Kept separately from ASSUMPTION_INFER_REGISTRY so that keys created *after* +# the rule is declared still receive it -- see AssumptionKey.__post_init__. +UNIVERSAL_RULES: list[tuple[tuple[type, ...], InferFactFn]] = [] + + +def _install_rule( + key: AssumptionKey, op_types: tuple[type, ...], fn: InferFactFn +) -> None: + for op_type in op_types: + ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), []).append(fn) + + +def register_universal_assumption( + *op_types: type, +) -> Callable[[InferFactFn], InferFactFn]: + """Decorator registering an inference rule that applies to *every* assumption key. + + Use this for rules that are indifferent to what the property means -- an Op that + forwards its input unchanged, or one that delegates to another Op. The rule is + installed for keys that already exist and for every key created later. + + Parameters + ---------- + *op_types : type + Op classes the rule applies to. + """ + + def decorator(fn: InferFactFn) -> InferFactFn: + UNIVERSAL_RULES.append((op_types, fn)) + for key in KEY_REGISTRY.values(): + _install_rule(key, op_types, fn) + return fn + + return decorator + + +class KeyRegistryView(Sequence[AssumptionKey]): + """Live, read-only sequence of every registered :class:`AssumptionKey`.""" + + __slots__ = () + + def __getitem__(self, index: int) -> AssumptionKey: + return tuple(KEY_REGISTRY.values())[index] + + def __iter__(self) -> Iterator[AssumptionKey]: + # Snapshot: a rule that constructs a key would otherwise resize the registry + # mid-iteration. Defining this at all keeps iteration off the inherited + # ``Sequence.__iter__``, which walks ``__getitem__`` index by index. + return iter(tuple(KEY_REGISTRY.values())) + + def __contains__(self, value) -> bool: + return KEY_REGISTRY.get(getattr(value, "name", None)) == value + + def __len__(self) -> int: + return len(KEY_REGISTRY) + + def __repr__(self) -> str: + return f"({', '.join(KEY_REGISTRY)})" + + # 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]] = {} @@ -116,7 +210,10 @@ def register_constant_inference(key: AssumptionKey, fn: ConstantInferFn) -> None PERMUTATION, ) -ALL_KEYS = (*MATRIX_KEYS, UNIQUE_INDICES) +# Live view rather than a tuple: a key registered by a downstream library shows up here +# too, so anything that iterates the keys at call time (debugprint, the drain rewrite) +# covers it without further registration. +ALL_KEYS = KeyRegistryView() # Implications about structural properties derivably from other structural properties register_implies(DIAGONAL, LOWER_TRIANGULAR, UPPER_TRIANGULAR, SYMMETRIC) diff --git a/pytensor/assumptions/specify.py b/pytensor/assumptions/specify.py index bb9f77d6ae..69e62727d5 100644 --- a/pytensor/assumptions/specify.py +++ b/pytensor/assumptions/specify.py @@ -1,6 +1,6 @@ from collections.abc import Sequence -from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption +from pytensor.assumptions.core import FactState, register_universal_assumption from pytensor.compile.ops import TypeCastingOp from pytensor.graph.basic import Apply, Variable from pytensor.tensor import TensorLike @@ -48,6 +48,7 @@ def pullback( return list(output_cotangents) +@register_universal_assumption(SpecifyAssumptions) def specify_assumption_rule(key, op, feature, fgraph, node, input_states): """Report the declared state for ``key`` joined with whatever inference derived from the input. The join surfaces ``ConflictingAssumptionsError`` when the user @@ -103,7 +104,6 @@ def assume( aliases a non-negative one (e.g. ``-1`` and ``n-1``). Such an index can never enlarge the axis it indexes, so it can be lifted earlier through operations without risk of duplicating computation. - Returns ------- out : TensorVariable @@ -111,10 +111,13 @@ def assume( Examples -------- - >>> import pytensor.tensor as pt - >>> x = pt.dmatrix("x") - >>> x_diag = assume(x, diagonal=True) - >>> x_not_sym = assume(x, symmetric=False) + .. code-block:: python + + import pytensor.tensor as pt + + x = pt.dmatrix("x") + x_diag = assume(x, diagonal=True) + x_not_sym = assume(x, symmetric=False) """ if not isinstance(x, Variable): x = as_tensor_variable(x) @@ -130,17 +133,13 @@ def assume( "permutation": permutation, "unique_indices": unique_indices, } - assumptions = { + declared = { name: FactState.TRUE if value else FactState.FALSE for name, value in values.items() if value is not None } - if not assumptions: + if not declared: return x - return SpecifyAssumptions(assumptions)(x) - - -for _key in ALL_KEYS: - register_assumption(_key, SpecifyAssumptions)(specify_assumption_rule) + return SpecifyAssumptions(declared)(x) diff --git a/pytensor/tensor/rewriting/assumptions.py b/pytensor/tensor/rewriting/assumptions.py index a60b7a798d..47046e6e7e 100644 --- a/pytensor/tensor/rewriting/assumptions.py +++ b/pytensor/tensor/rewriting/assumptions.py @@ -1,12 +1,9 @@ -from pytensor.assumptions import ALL_KEYS, AssumptionFeature +from pytensor.assumptions import KEY_REGISTRY, AssumptionFeature from pytensor.assumptions.specify import SpecifyAssumptions from pytensor.compile.mode import optdb from pytensor.graph.rewriting.basic import GraphRewriter -_KEY_BY_NAME = {key.name: key for key in ALL_KEYS} - - class DrainSpecifyAssumptions(GraphRewriter): """Drain ``SpecifyAssumptions`` declarations into the ``AssumptionFeature`` and remove the marker nodes. @@ -39,15 +36,14 @@ def apply(self, fgraph): [out] = node.outputs # Resolve the asserted facts into the cache. for name, _ in node.op.assumptions: - assumption_feature.get(out, _KEY_BY_NAME[name]) - # Drain the marker: redirect its consumers to the raw input, - # peeling nested SpecifyAssumptions so a single replace_all - # collapses ``assume(assume(...))`` chains all the way down. + assumption_feature.get(out, KEY_REGISTRY[name]) + # Drain the marker: redirect its consumers to the raw input, peeling + # already-drained nested SpecifyAssumptions -- ``nodes`` is in toposort + # order, so a single replace_all collapses ``assume(assume(...))`` chains + # all the way down. inp = node.inputs[0] - while inp.owner is not None and isinstance( - inp.owner.op, SpecifyAssumptions - ): - inp = inp.owner.inputs[0] + while inp in replacements: + inp = replacements[inp] replacements[out] = inp fgraph.replace_all( diff --git a/tests/assumptions/conftest.py b/tests/assumptions/conftest.py index b82e5d10c3..a8bd3d1990 100644 --- a/tests/assumptions/conftest.py +++ b/tests/assumptions/conftest.py @@ -12,6 +12,8 @@ def _snapshot_assumption_registries(): """Restore module-global assumption registries after each test.""" infer_snapshot = copy.deepcopy(_assumptions_core.ASSUMPTION_INFER_REGISTRY) implies_snapshot = copy.deepcopy(_assumptions_core.IMPLIES) + key_snapshot = dict(_assumptions_core.KEY_REGISTRY) + universal_snapshot = list(_assumptions_core.UNIVERSAL_RULES) try: yield finally: @@ -19,6 +21,9 @@ def _snapshot_assumption_registries(): _assumptions_core.ASSUMPTION_INFER_REGISTRY.update(infer_snapshot) _assumptions_core.IMPLIES.clear() _assumptions_core.IMPLIES.update(implies_snapshot) + _assumptions_core.KEY_REGISTRY.clear() + _assumptions_core.KEY_REGISTRY.update(key_snapshot) + _assumptions_core.UNIVERSAL_RULES[:] = universal_snapshot def make_fgraph(*outputs, **kwargs): From 5beba2b6b00cc04f20522dabad8bb9b806b43ebb Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 22 Aug 2026 22:45:37 -0400 Subject: [PATCH 2/9] Let downstream libraries declare assumptions assume() keeps its named arguments for the built-in properties so they stay documented and typo-checked, and accepts registered extension keys through **assumptions. --- pytensor/assumptions/core.py | 67 +++++++++++- pytensor/assumptions/specify.py | 25 ++++- tests/assumptions/test_extension.py | 154 ++++++++++++++++++++++++++++ 3 files changed, 239 insertions(+), 7 deletions(-) create mode 100644 tests/assumptions/test_extension.py diff --git a/pytensor/assumptions/core.py b/pytensor/assumptions/core.py index 39ce5d9899..e57d4d6994 100644 --- a/pytensor/assumptions/core.py +++ b/pytensor/assumptions/core.py @@ -5,8 +5,10 @@ from typing import Any from pytensor.graph import Apply, FunctionGraph, Op +from pytensor.graph.basic import Variable from pytensor.graph.features import AlreadyThere, Feature from pytensor.graph.traversal import walk_toposort +from pytensor.tensor import TensorLike from pytensor.tensor.variable import TensorConstant @@ -45,16 +47,29 @@ class AssumptionKey: declared with :func:`register_universal_assumption` for it, which is all a downstream library must do to add a property of its own. - ``name`` is the key's identity: two keys may not share one unless they are - identical in every field. + ``name`` is the key's identity: it is what :func:`assume` accepts as a keyword and + what two keys may not share unless they are identical in every field. Parameters ---------- name : str - Unique identifier for the property. + Unique identifier for the property, also the keyword :func:`assume` accepts. short_name : str, optional Abbreviated label used by ``debugprint(print_assumptions=True)``. Falls back to ``name`` when empty. + + Examples + -------- + .. code-block:: python + + import pytensor.tensor as pt + from pytensor.assumptions import AssumptionKey + + TOEPLITZ = AssumptionKey("toeplitz", short_name="toep") + + x = TOEPLITZ.assume(pt.matrix("x")) + TOEPLITZ.holds(x) # True + TOEPLITZ.holds(pt.exp(x)) # False -- no rule teaches it about Elemwise yet """ name: str @@ -80,6 +95,49 @@ def __post_init__(self) -> None: def __repr__(self) -> str: return self.name + def assume(self, x: TensorLike, *, state: bool = True) -> Variable: + """Return a view of *x* declaring this assumption. + + Parameters + ---------- + x : tensor-like + The input to annotate. + state : bool, optional + Whether to assert the property holds or that it does not. Default True. + + Returns + ------- + out : TensorVariable + A view of *x* with the assumption attached. + """ + from pytensor.assumptions.specify import SpecifyAssumptions + + fact = FactState.TRUE if state else FactState.FALSE + return SpecifyAssumptions({self.name: fact})(x) + + def holds(self, var: Variable, fgraph: FunctionGraph | None = None) -> bool: + """Return True iff this assumption is provably TRUE for *var*. + + Parameters + ---------- + var : Variable + The variable to ask about. + fgraph : FunctionGraph, optional + Graph to resolve the question in. Pass one when asking about several + variables of the same graph: without it a throwaway ``FunctionGraph`` is + built per call, which walks the ancestors of *var* and discards the + inference cache afterwards. + + Returns + ------- + holds : bool + True when the property is known to hold, False when it is known not to + hold or is simply unknown. + """ + if fgraph is None: + fgraph = FunctionGraph(outputs=[var], clone=False) + return check_assumption(fgraph, var, self) + class ConflictingAssumptionsError(ValueError): """Raised when joining evidence about a (variable, key) produces ``FactState.CONFLICT``. @@ -102,7 +160,8 @@ class ConflictingAssumptionsError(ValueError): ASSUMPTION_INFER_REGISTRY: dict[tuple[AssumptionKey, type], list[InferFactFn]] = {} # Every AssumptionKey ever constructed, by name. Downstream libraries join the system -# by constructing a key; nothing else is required of them. +# by constructing a key; nothing else is required of them. This is what resolves the +# names :func:`assume` takes as keywords. KEY_REGISTRY: dict[str, AssumptionKey] = {} # Rules that hold for every key regardless of what the key means, as (op_types, fn) diff --git a/pytensor/assumptions/specify.py b/pytensor/assumptions/specify.py index 69e62727d5..d07c96225b 100644 --- a/pytensor/assumptions/specify.py +++ b/pytensor/assumptions/specify.py @@ -1,6 +1,10 @@ from collections.abc import Sequence -from pytensor.assumptions.core import FactState, register_universal_assumption +from pytensor.assumptions.core import ( + KEY_REGISTRY, + FactState, + register_universal_assumption, +) from pytensor.compile.ops import TypeCastingOp from pytensor.graph.basic import Apply, Variable from pytensor.tensor import TensorLike @@ -71,6 +75,7 @@ def assume( selection: bool | None = None, permutation: bool | None = None, unique_indices: bool | None = None, + **assumptions: bool | None, ): """Attach structural assumptions to a symbolic tensor. @@ -104,6 +109,10 @@ def assume( aliases a non-negative one (e.g. ``-1`` and ``n-1``). Such an index can never enlarge the axis it indexes, so it can be lifted earlier through operations without risk of duplicating computation. + **assumptions : bool, optional + Assumptions registered by downstream libraries, passed by key name, e.g. + ``time_varying=True``. + Returns ------- out : TensorVariable @@ -122,7 +131,7 @@ def assume( if not isinstance(x, Variable): x = as_tensor_variable(x) - values = { + core_values = { "diagonal": diagonal, "lower_triangular": lower_triangular, "upper_triangular": upper_triangular, @@ -133,9 +142,19 @@ def assume( "permutation": permutation, "unique_indices": unique_indices, } + + unknown = [name for name in assumptions if name not in KEY_REGISTRY] + if unknown: + extensions = sorted(KEY_REGISTRY.keys() - core_values.keys()) + raise ValueError( + f"Unknown assumption(s): {', '.join(unknown)}. Registered extension " + f"assumptions are: {', '.join(extensions) if extensions else '(none)'}. " + f"Register a new one by constructing an AssumptionKey." + ) + declared = { name: FactState.TRUE if value else FactState.FALSE - for name, value in values.items() + for name, value in (core_values | assumptions).items() if value is not None } diff --git a/tests/assumptions/test_extension.py b/tests/assumptions/test_extension.py new file mode 100644 index 0000000000..e4269a1a78 --- /dev/null +++ b/tests/assumptions/test_extension.py @@ -0,0 +1,154 @@ +import pytest + +import pytensor.tensor as pt +from pytensor.assumptions import ( + ALL_KEYS, + KEY_REGISTRY, + SYMMETRIC, + AssumptionKey, + FactState, + register_assumption, + register_universal_assumption, +) +from pytensor.assumptions.specify import SpecifyAssumptions, assume +from pytensor.printing import debugprint +from pytensor.tensor.basic import AllocDiag, alloc_diag +from pytensor.tensor.blockwise import Blockwise +from pytensor.tensor.rewriting.assumptions import DrainSpecifyAssumptions +from tests.assumptions.conftest import make_fgraph + + +def test_key_registers_itself(): + key = AssumptionKey("time_varying", short_name="tv") + assert KEY_REGISTRY["time_varying"] is key + assert key in ALL_KEYS + + +def test_identical_key_redefinition_is_idempotent(): + """A module imported twice must not double-register its key or its rules.""" + first = AssumptionKey("time_varying", short_name="tv") + n_keys = len(KEY_REGISTRY) + second = AssumptionKey("time_varying", short_name="tv") + + assert second == first + assert len(KEY_REGISTRY) == n_keys + assert KEY_REGISTRY["time_varying"] is first + + +def test_name_collision_with_different_metadata_raises(): + AssumptionKey("time_varying", short_name="tv") + with pytest.raises(ValueError, match="already registered"): + AssumptionKey("time_varying", short_name="clashing") + + +def test_assume_accepts_extension_key(): + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x = pt.tensor3("x") + x_tv = assume(x, time_varying=True) + _, af = make_fgraph(x_tv) + assert af.check(x_tv, TIME_VARYING) + + +def test_assume_mixes_core_and_extension_keys(): + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x = pt.tensor("x", shape=(10, 3, 3)) + x_both = assume(x, symmetric=True, time_varying=True) + _, af = make_fgraph(x_both) + assert af.check(x_both, SYMMETRIC) + assert af.check(x_both, TIME_VARYING) + + +def test_assume_records_false_for_extension_key(): + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x = pt.tensor3("x") + x_static = assume(x, time_varying=False) + _, af = make_fgraph(x_static) + assert af.get(x_static, TIME_VARYING) is FactState.FALSE + + +def test_assume_rejects_unregistered_name(): + """A typo must not silently become a no-op, and the error must aid discovery.""" + AssumptionKey("time_varying", short_name="tv") + x = pt.matrix("x") + with pytest.raises(ValueError, match="Unknown assumption\\(s\\): symmetrik"): + assume(x, symmetrik=True) + with pytest.raises(ValueError, match="are: time_varying"): + assume(x, symmetrik=True) + + +def test_key_assume_and_holds(): + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x = pt.tensor3("x") + + assert TIME_VARYING.holds(TIME_VARYING.assume(x)) + assert not TIME_VARYING.holds(TIME_VARYING.assume(x, state=False)) + assert not TIME_VARYING.holds(x) + + +def test_holds_reuses_a_supplied_fgraph(): + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x_tv = TIME_VARYING.assume(pt.tensor3("x")) + y = pt.matrix("y") + fgraph, _ = make_fgraph(x_tv, y) + + assert TIME_VARYING.holds(x_tv, fgraph) + assert not TIME_VARYING.holds(y, fgraph) + + +def test_universal_rules_reach_a_later_key(): + """A key created after ``blockwise`` was imported still gets its delegate.""" + SPARSE = AssumptionKey("sparse") + register_assumption(SPARSE, AllocDiag)( + lambda key, op, feature, fgraph, node, input_states: [FactState.TRUE] + ) + + v_core = pt.vector("v", shape=(3,)) + core_op = alloc_diag(v_core, offset=0, axis1=0, axis2=1).owner.op + v_batch = pt.matrix("v_batch", shape=(5, 3)) + batched = Blockwise(core_op, signature="(n)->(n,n)")(v_batch) + + _, af = make_fgraph(batched) + assert af.check(batched, SPARSE) + + +def test_universal_rule_reaches_existing_keys(): + """The decorator installs onto keys registered before it ran, not just after.""" + EARLY = AssumptionKey("early") + + @register_universal_assumption(AllocDiag) + def _always_true(key, op, feature, fgraph, node, input_states): + return [FactState.TRUE] + + LATE = AssumptionKey("late") + + diag = alloc_diag(pt.vector("v", shape=(3,)), offset=0, axis1=0, axis2=1) + _, af = make_fgraph(diag) + + assert af.check(diag, EARLY) + assert af.check(diag, LATE) + + +def test_membership_rejects_a_non_key_sharing_a_name(): + """``in`` compares keys, not names -- a variable named after one is not a key.""" + assert "symmetric" not in ALL_KEYS + assert pt.matrix("symmetric") not in ALL_KEYS + + +def test_extension_key_appears_in_debugprint(): + AssumptionKey("time_varying", short_name="tv") + x_tv = assume(pt.tensor3("x"), time_varying=True) + printed = debugprint(x_tv, print_assumptions=True, file="str") + assert "a={tv}" in printed + + +def test_drain_resolves_extension_key(): + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x = pt.tensor3("x") + fgraph, af = make_fgraph(assume(x, time_varying=True) + 1, inputs=[x]) + + DrainSpecifyAssumptions().apply(fgraph) + + assert not any( + isinstance(node.op, SpecifyAssumptions) for node in fgraph.apply_nodes + ) + assert af.check(x, TIME_VARYING) From 6bc7effe0421a2eeded85d1caa51da23fe782abb Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 22 Aug 2026 22:46:01 -0400 Subject: [PATCH 3/9] Key SpecifyAssumptions by AssumptionKey instead of by name Holding a key implies it is registered, so a graph can no longer carry a declaration the system cannot resolve and the drain rewrite needs no lookup. AssumptionKey.__reduce__ keeps that true across pickling, where the default dataclass path would skip __init__ and restore a key with no rules installed. --- pytensor/assumptions/core.py | 13 ++++++-- pytensor/assumptions/specify.py | 41 ++++++++++++++++------- pytensor/tensor/rewriting/assumptions.py | 6 ++-- tests/assumptions/test_extension.py | 42 ++++++++++++++++++++++++ tests/assumptions/test_specify.py | 4 +-- 5 files changed, 86 insertions(+), 20 deletions(-) diff --git a/pytensor/assumptions/core.py b/pytensor/assumptions/core.py index e57d4d6994..f265786db3 100644 --- a/pytensor/assumptions/core.py +++ b/pytensor/assumptions/core.py @@ -92,6 +92,12 @@ def __post_init__(self) -> None: for op_types, fn in UNIVERSAL_RULES: _install_rule(self, op_types, fn) + def __reduce__(self): + # Unpickle through the constructor: the default dataclass path skips __init__, + # leaving a key that is in no registry and has no rules installed -- not even + # the one that reads declarations back off SpecifyAssumptions. + return type(self), (self.name, self.short_name) + def __repr__(self) -> str: return self.name @@ -113,7 +119,7 @@ def assume(self, x: TensorLike, *, state: bool = True) -> Variable: from pytensor.assumptions.specify import SpecifyAssumptions fact = FactState.TRUE if state else FactState.FALSE - return SpecifyAssumptions({self.name: fact})(x) + return SpecifyAssumptions({self: fact})(x) def holds(self, var: Variable, fgraph: FunctionGraph | None = None) -> bool: """Return True iff this assumption is provably TRUE for *var*. @@ -160,8 +166,9 @@ class ConflictingAssumptionsError(ValueError): ASSUMPTION_INFER_REGISTRY: dict[tuple[AssumptionKey, type], list[InferFactFn]] = {} # Every AssumptionKey ever constructed, by name. Downstream libraries join the system -# by constructing a key; nothing else is required of them. This is what resolves the -# names :func:`assume` takes as keywords. +# by constructing a key; nothing else is required of them. This resolves the names +# :func:`assume` takes as keywords -- graphs themselves carry keys, not names, so +# nothing downstream of graph construction needs to look anything up here. KEY_REGISTRY: dict[str, AssumptionKey] = {} # Rules that hold for every key regardless of what the key means, as (op_types, fn) diff --git a/pytensor/assumptions/specify.py b/pytensor/assumptions/specify.py index d07c96225b..a79ce371ff 100644 --- a/pytensor/assumptions/specify.py +++ b/pytensor/assumptions/specify.py @@ -2,6 +2,7 @@ from pytensor.assumptions.core import ( KEY_REGISTRY, + AssumptionKey, FactState, register_universal_assumption, ) @@ -14,26 +15,42 @@ class SpecifyAssumptions(TypeCastingOp): """No-op that declares structural assumptions on a tensor for use by graph rewrites. - ``assumptions`` is a tuple of ``(name, FactState)`` pairs sorted by ``name``, where - ``name`` matches the name of an :class:`AssumptionKey`. Two instances with the same - fact set compare equal via ``__props__``, so PyTensor's graph merge collapses - duplicates. + ``assumptions`` is a tuple of ``(AssumptionKey, FactState)`` pairs sorted by key + name. Declaring a fact therefore requires holding the key itself, and constructing + a key registers it, so a graph cannot carry an assumption the system has never + heard of. Two instances with the same fact set compare equal via ``__props__``, so + PyTensor's graph merge collapses duplicates. + + Parameters + ---------- + assumptions : dict mapping AssumptionKey to FactState + The facts to declare. """ __props__ = ("assumptions",) - assumptions: tuple[tuple[str, FactState], ...] + assumptions: tuple[tuple[AssumptionKey, FactState], ...] - def __init__(self, assumptions: dict[str, FactState]): + def __init__(self, assumptions: dict[AssumptionKey, FactState]): super().__init__() + passed_by_name = [ + key for key in assumptions if not isinstance(key, AssumptionKey) + ] + if passed_by_name: + raise TypeError( + f"SpecifyAssumptions is keyed by AssumptionKey, not by name: " + f"{passed_by_name!r}. Pass the key objects, or declare by name " + f"with assume()." + ) self.assumptions = tuple( - (name, FactState(state)) for name, state in sorted(assumptions.items()) + (key, FactState(state)) + for key, state in sorted(assumptions.items(), key=lambda kv: kv[0].name) ) def __str__(self): facts = ", ".join( - name if state is FactState.TRUE else f"!{name}" - for name, state in self.assumptions + key.name if state is FactState.TRUE else f"!{key.name}" + for key, state in self.assumptions ) return f"{type(self).__name__}{{{facts}}}" @@ -58,8 +75,8 @@ def specify_assumption_rule(key, op, feature, fgraph, node, input_states): from the input. The join surfaces ``ConflictingAssumptionsError`` when the user asserts a state that contradicts what the system can prove (e.g. asserting ``diagonal=False`` on something proved diagonal).""" - for name, state in op.assumptions: - if name == key.name: + for declared_key, state in op.assumptions: + if declared_key == key: return [FactState.join(state, input_states[0])] return [input_states[0]] @@ -153,7 +170,7 @@ def assume( ) declared = { - name: FactState.TRUE if value else FactState.FALSE + KEY_REGISTRY[name]: FactState.TRUE if value else FactState.FALSE for name, value in (core_values | assumptions).items() if value is not None } diff --git a/pytensor/tensor/rewriting/assumptions.py b/pytensor/tensor/rewriting/assumptions.py index 47046e6e7e..3740863eef 100644 --- a/pytensor/tensor/rewriting/assumptions.py +++ b/pytensor/tensor/rewriting/assumptions.py @@ -1,4 +1,4 @@ -from pytensor.assumptions import KEY_REGISTRY, AssumptionFeature +from pytensor.assumptions import AssumptionFeature from pytensor.assumptions.specify import SpecifyAssumptions from pytensor.compile.mode import optdb from pytensor.graph.rewriting.basic import GraphRewriter @@ -35,8 +35,8 @@ def apply(self, fgraph): for node in nodes: [out] = node.outputs # Resolve the asserted facts into the cache. - for name, _ in node.op.assumptions: - assumption_feature.get(out, KEY_REGISTRY[name]) + for key, _ in node.op.assumptions: + assumption_feature.get(out, key) # Drain the marker: redirect its consumers to the raw input, peeling # already-drained nested SpecifyAssumptions -- ``nodes`` is in toposort # order, so a single replace_all collapses ``assume(assume(...))`` chains diff --git a/tests/assumptions/test_extension.py b/tests/assumptions/test_extension.py index e4269a1a78..3c07a4dd71 100644 --- a/tests/assumptions/test_extension.py +++ b/tests/assumptions/test_extension.py @@ -1,3 +1,5 @@ +import pickle + import pytest import pytensor.tensor as pt @@ -10,6 +12,7 @@ register_assumption, register_universal_assumption, ) +from pytensor.assumptions.core import ASSUMPTION_INFER_REGISTRY from pytensor.assumptions.specify import SpecifyAssumptions, assume from pytensor.printing import debugprint from pytensor.tensor.basic import AllocDiag, alloc_diag @@ -152,3 +155,42 @@ def test_drain_resolves_extension_key(): isinstance(node.op, SpecifyAssumptions) for node in fgraph.apply_nodes ) assert af.check(x, TIME_VARYING) + + +def test_graph_carries_keys_not_names(): + """The declaration holds the key itself, so it cannot name an unregistered fact.""" + TIME_VARYING = AssumptionKey("time_varying", short_name="tv") + x_tv = assume(pt.tensor3("x"), time_varying=True) + assert x_tv.owner.op.assumptions == ((TIME_VARYING, FactState.TRUE),) + + +def test_declaring_by_name_is_rejected(): + AssumptionKey("time_varying", short_name="tv") + with pytest.raises(TypeError, match="not by name"): + SpecifyAssumptions({"time_varying": FactState.TRUE}) + + +def test_key_survives_a_pickle_round_trip(): + """A key restored from a cached graph re-registers and keeps its universal rules.""" + key = AssumptionKey("time_varying", short_name="tv") + blob = pickle.dumps(key) + del KEY_REGISTRY["time_varying"] + + restored = pickle.loads(blob) + + assert KEY_REGISTRY["time_varying"] == restored + assert (restored, SpecifyAssumptions) in ASSUMPTION_INFER_REGISTRY + assert restored.holds(restored.assume(pt.tensor3("x"))) + + +def test_pickled_graph_keeps_its_declaration(): + """A graph outliving the library that declared its key still drains the fact.""" + key = AssumptionKey("time_varying", short_name="tv") + blob = pickle.dumps(key.assume(pt.tensor3("x"))) + del KEY_REGISTRY["time_varying"] + + restored_graph = pickle.loads(blob) + [(restored_key, state)] = restored_graph.owner.op.assumptions + + assert state is FactState.TRUE + assert restored_key.holds(restored_graph) diff --git a/tests/assumptions/test_specify.py b/tests/assumptions/test_specify.py index 0847f48bd5..3a2a8ebd6d 100644 --- a/tests/assumptions/test_specify.py +++ b/tests/assumptions/test_specify.py @@ -68,8 +68,8 @@ def test_assume_chained_combines_facts(): def test_specify_assumptions_op_equal_for_same_facts(): - a = SpecifyAssumptions({"diagonal": FactState.TRUE, "symmetric": FactState.FALSE}) - b = SpecifyAssumptions({"symmetric": FactState.FALSE, "diagonal": FactState.TRUE}) + a = SpecifyAssumptions({DIAGONAL: FactState.TRUE, SYMMETRIC: FactState.FALSE}) + b = SpecifyAssumptions({SYMMETRIC: FactState.FALSE, DIAGONAL: FactState.TRUE}) assert a == b assert hash(a) == hash(b) From ce324286d4372860a49dd9c9b8d885d3c4841e70 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 22 Aug 2026 23:17:55 -0400 Subject: [PATCH 4/9] Stop propagating matrix properties to unique_indices The trailing-two-axes rules were registered for every key, so Alloc broadcasting an index array claimed the result still had distinct entries while duplicating every one of them -- a fact that exists to license lifting an index earlier. --- pytensor/assumptions/alloc.py | 4 ++-- pytensor/assumptions/dimshuffle.py | 4 ++-- pytensor/assumptions/reshape.py | 4 ++-- pytensor/assumptions/shape.py | 4 ++-- pytensor/assumptions/subtensor.py | 4 ++-- tests/assumptions/test_alloc.py | 17 +++++++++++++++++ 6 files changed, 27 insertions(+), 10 deletions(-) diff --git a/pytensor/assumptions/alloc.py b/pytensor/assumptions/alloc.py index a84c84ce83..35c498aea2 100644 --- a/pytensor/assumptions/alloc.py +++ b/pytensor/assumptions/alloc.py @@ -1,5 +1,5 @@ from pytensor.assumptions.core import ( - ALL_KEYS, + MATRIX_KEYS, FactState, register_assumption, true_if, @@ -124,5 +124,5 @@ def alloc_propagates_matrix_property( return [FactState.UNKNOWN] -for _key in ALL_KEYS: +for _key in MATRIX_KEYS: register_assumption(_key, Alloc)(alloc_propagates_matrix_property) diff --git a/pytensor/assumptions/dimshuffle.py b/pytensor/assumptions/dimshuffle.py index e60edcd219..0f622e73a1 100644 --- a/pytensor/assumptions/dimshuffle.py +++ b/pytensor/assumptions/dimshuffle.py @@ -1,4 +1,4 @@ -from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption +from pytensor.assumptions.core import MATRIX_KEYS, FactState, register_assumption from pytensor.tensor.elemwise import DimShuffle @@ -29,5 +29,5 @@ def dimshuffle_propagates_matrix_property( return [FactState.UNKNOWN] -for _key in ALL_KEYS: +for _key in MATRIX_KEYS: register_assumption(_key, DimShuffle)(dimshuffle_propagates_matrix_property) diff --git a/pytensor/assumptions/reshape.py b/pytensor/assumptions/reshape.py index 53ec7d3a57..41a5f09a34 100644 --- a/pytensor/assumptions/reshape.py +++ b/pytensor/assumptions/reshape.py @@ -1,4 +1,4 @@ -from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption +from pytensor.assumptions.core import MATRIX_KEYS, FactState, register_assumption from pytensor.tensor.reshape import JoinDims, SplitDims @@ -23,6 +23,6 @@ def split_dims_propagates_matrix_property( return [FactState.UNKNOWN] -for _key in ALL_KEYS: +for _key in MATRIX_KEYS: register_assumption(_key, JoinDims)(join_dims_propagates_matrix_property) register_assumption(_key, SplitDims)(split_dims_propagates_matrix_property) diff --git a/pytensor/assumptions/shape.py b/pytensor/assumptions/shape.py index 1e907d94a5..a036de18b3 100644 --- a/pytensor/assumptions/shape.py +++ b/pytensor/assumptions/shape.py @@ -1,4 +1,4 @@ -from pytensor.assumptions.core import ALL_KEYS, FactState, register_assumption +from pytensor.assumptions.core import MATRIX_KEYS, FactState, register_assumption from pytensor.tensor.shape import Reshape, SpecifyShape @@ -25,6 +25,6 @@ def reshape_propagates_matrix_property( return [input_states[0]] -for _key in ALL_KEYS: +for _key in MATRIX_KEYS: register_assumption(_key, SpecifyShape)(specify_shape_propagates_matrix_property) register_assumption(_key, Reshape)(reshape_propagates_matrix_property) diff --git a/pytensor/assumptions/subtensor.py b/pytensor/assumptions/subtensor.py index 8f13338676..b7ec162c1b 100644 --- a/pytensor/assumptions/subtensor.py +++ b/pytensor/assumptions/subtensor.py @@ -1,7 +1,7 @@ from pytensor.assumptions.core import ( - ALL_KEYS, DIAGONAL, LOWER_TRIANGULAR, + MATRIX_KEYS, POSITIVE_DEFINITE, SYMMETRIC, UPPER_TRIANGULAR, @@ -82,5 +82,5 @@ def incsubtensor_propagates_matrix_property( return true_if(base_state is FactState.TRUE and value_state is FactState.TRUE) -for _key in ALL_KEYS: +for _key in MATRIX_KEYS: register_assumption(_key, IncSubtensor)(incsubtensor_propagates_matrix_property) diff --git a/tests/assumptions/test_alloc.py b/tests/assumptions/test_alloc.py index c248770800..596ded986b 100644 --- a/tests/assumptions/test_alloc.py +++ b/tests/assumptions/test_alloc.py @@ -10,6 +10,7 @@ PERMUTATION, POSITIVE_DEFINITE, SYMMETRIC, + UNIQUE_INDICES, UPPER_TRIANGULAR, FactState, ) @@ -145,3 +146,19 @@ def test_alloc_broadcast_vector_value_is_unknown(): y = pt.alloc(v, 4, 4) _, af = make_fgraph(y) assert af.get(y, SYMMETRIC) == FactState.UNKNOWN + + +def test_unique_indices_survives_no_broadcast(): + """Alloc repeats entries, so a uniqueness claim must not carry through it. + + The matrix-property rules propagate anything whose trailing two axes are untouched, + which is wrong for a claim about the values themselves. + """ + idx = pt.matrix("idx", shape=(2, 3), dtype="int64") + broadcast = pt.alloc(assume(idx, unique_indices=True), 4, 2, 3) + + _, af = make_fgraph(broadcast) + assert af.get(broadcast, UNIQUE_INDICES) is not FactState.TRUE + + repeated = broadcast.eval({idx: np.arange(6).reshape(2, 3)}) + assert len(np.unique(repeated)) < repeated.size From a68b1f88084620abfb3743ee7573df0375f1c7a4 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 22 Aug 2026 23:18:16 -0400 Subject: [PATCH 5/9] Add register_matrix_property_rules for downstream matrix properties A key wanting all but one of the rules registers its own with prepend=True, since rules are tried in registration order until one answers. --- pytensor/assumptions/__init__.py | 1 + pytensor/assumptions/bundles.py | 55 ++++++++++++++++++++ pytensor/assumptions/core.py | 19 ++++++- tests/assumptions/test_bundles.py | 86 +++++++++++++++++++++++++++++++ 4 files changed, 159 insertions(+), 2 deletions(-) create mode 100644 pytensor/assumptions/bundles.py create mode 100644 tests/assumptions/test_bundles.py diff --git a/pytensor/assumptions/__init__.py b/pytensor/assumptions/__init__.py index 8bfe7c37aa..1739b69bfc 100644 --- a/pytensor/assumptions/__init__.py +++ b/pytensor/assumptions/__init__.py @@ -14,6 +14,7 @@ import pytensor.assumptions.subtensor import pytensor.assumptions.symmetric import pytensor.assumptions.triangular +from pytensor.assumptions.bundles import register_matrix_property_rules from pytensor.assumptions.core import ( ALL_KEYS, DIAGONAL, diff --git a/pytensor/assumptions/bundles.py b/pytensor/assumptions/bundles.py new file mode 100644 index 0000000000..001eaf4bf5 --- /dev/null +++ b/pytensor/assumptions/bundles.py @@ -0,0 +1,55 @@ +from pytensor.assumptions.alloc import alloc_propagates_matrix_property +from pytensor.assumptions.core import AssumptionKey, register_assumption +from pytensor.assumptions.dimshuffle import dimshuffle_propagates_matrix_property +from pytensor.assumptions.reshape import ( + join_dims_propagates_matrix_property, + split_dims_propagates_matrix_property, +) +from pytensor.assumptions.shape import ( + reshape_propagates_matrix_property, + specify_shape_propagates_matrix_property, +) +from pytensor.assumptions.subtensor import ( + incsubtensor_propagates_matrix_property, + subtensor_propagates_matrix_property, +) +from pytensor.tensor.basic import Alloc +from pytensor.tensor.elemwise import DimShuffle +from pytensor.tensor.reshape import JoinDims, SplitDims +from pytensor.tensor.shape import Reshape, SpecifyShape +from pytensor.tensor.subtensor import IncSubtensor, Subtensor + + +def register_matrix_property_rules(key: AssumptionKey) -> None: + """Register the standard propagation rules for a property of the trailing two axes. + + Every rule here answers one question: does the Op leave the trailing two axes + undisturbed? The bundle thus suits any property of a matrix that batch dimensions + carry elementwise, such as triangularity or a fixed sparsity pattern. + + Rules are tried in registration order until one returns a non-UNKNOWN state, so a + key needing different behavior for one Op registers its own with + ``register_assumption(..., prepend=True)``. + + Parameters + ---------- + key : AssumptionKey + The property to install the rules for. + + Examples + -------- + .. code-block:: python + + from pytensor.assumptions import AssumptionKey, register_matrix_property_rules + + TOEPLITZ = AssumptionKey("toeplitz", short_name="toep") + register_matrix_property_rules(TOEPLITZ) + """ + register_assumption(key, DimShuffle)(dimshuffle_propagates_matrix_property) + register_assumption(key, Reshape)(reshape_propagates_matrix_property) + register_assumption(key, SpecifyShape)(specify_shape_propagates_matrix_property) + register_assumption(key, JoinDims)(join_dims_propagates_matrix_property) + register_assumption(key, SplitDims)(split_dims_propagates_matrix_property) + register_assumption(key, Alloc)(alloc_propagates_matrix_property) + register_assumption(key, Subtensor)(subtensor_propagates_matrix_property) + register_assumption(key, IncSubtensor)(incsubtensor_propagates_matrix_property) diff --git a/pytensor/assumptions/core.py b/pytensor/assumptions/core.py index f265786db3..3b0d1ed8c4 100644 --- a/pytensor/assumptions/core.py +++ b/pytensor/assumptions/core.py @@ -288,17 +288,32 @@ def register_constant_inference(key: AssumptionKey, fn: ConstantInferFn) -> None def register_assumption( - key: AssumptionKey, *op_types: type + key: AssumptionKey, *op_types: type, prepend: bool = False ) -> 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. + + Parameters + ---------- + key : AssumptionKey + The property the rule infers. + *op_types : type + Op classes the rule applies to. + prepend : bool, optional + Run this rule ahead of those already registered for the same pair rather than + after them. Rules are tried in order until one returns a non-UNKNOWN state, so + this is how a key overrides a rule it inherited from a bundle. Default False. """ def decorator(fn: InferFactFn) -> InferFactFn: for op_type in op_types: - ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), []).append(fn) + rules = ASSUMPTION_INFER_REGISTRY.setdefault((key, op_type), []) + if prepend: + rules.insert(0, fn) + else: + rules.append(fn) return fn return decorator diff --git a/tests/assumptions/test_bundles.py b/tests/assumptions/test_bundles.py new file mode 100644 index 0000000000..9cd210a3cc --- /dev/null +++ b/tests/assumptions/test_bundles.py @@ -0,0 +1,86 @@ +import pytest + +import pytensor.tensor as pt +from pytensor.assumptions import ( + MATRIX_KEYS, + AssumptionKey, + FactState, + assume, + register_assumption, + register_matrix_property_rules, +) +from pytensor.assumptions.core import ASSUMPTION_INFER_REGISTRY +from pytensor.tensor.basic import Alloc +from pytensor.tensor.elemwise import DimShuffle +from pytensor.tensor.reshape import JoinDims, SplitDims +from pytensor.tensor.shape import Reshape, SpecifyShape +from pytensor.tensor.subtensor import IncSubtensor +from tests.assumptions.conftest import make_fgraph + + +def test_bundle_propagates_through_the_standard_ops(): + TOEPLITZ = AssumptionKey("toeplitz", short_name="toep") + register_matrix_property_rules(TOEPLITZ) + + x = assume(pt.tensor("x", shape=(4, 3, 3)), toeplitz=True) + + indexed = x[0] + reshaped = x.reshape((2, 2, 3, 3)) + shape_specified = pt.specify_shape(x, (4, 3, 3)) + + _, af = make_fgraph(indexed, reshaped, shape_specified) + assert af.check(indexed, TOEPLITZ) + assert af.check(reshaped, TOEPLITZ) + assert af.check(shape_specified, TOEPLITZ) + + +def test_bundle_leaves_the_core_axes_alone(): + """The rules protect the trailing two axes -- disturbing them stops propagation.""" + TOEPLITZ = AssumptionKey("toeplitz", short_name="toep") + register_matrix_property_rules(TOEPLITZ) + + x = assume(pt.matrix("x", shape=(3, 3)), toeplitz=True) + transposed = x.T + + _, af = make_fgraph(transposed) + assert af.get(transposed, TOEPLITZ) is FactState.UNKNOWN + + +@pytest.mark.parametrize( + "prepend, expected", + [(True, FactState.FALSE), (False, FactState.TRUE)], + ids=["prepend-wins", "append-loses"], +) +def test_prepend_decides_which_rule_answers(prepend, expected): + """Expand-dims is a case the bundle answers, so only order decides the outcome.""" + TOEPLITZ = AssumptionKey("toeplitz", short_name="toep") + register_matrix_property_rules(TOEPLITZ) + + @register_assumption(TOEPLITZ, DimShuffle, prepend=prepend) + def _never_survives_expand_dims(key, op, feature, fgraph, node, input_states): + return [FactState.FALSE] if op.is_expand_dims else [FactState.UNKNOWN] + + x = assume(pt.matrix("x", shape=(3, 3)), toeplitz=True) + expanded = x[None] + + _, af = make_fgraph(expanded) + assert af.get(expanded, TOEPLITZ) is expected + + +@pytest.mark.parametrize("key", MATRIX_KEYS, ids=lambda k: k.name) +@pytest.mark.parametrize( + "op_type", + [DimShuffle, Reshape, SpecifyShape, JoinDims, SplitDims, Alloc, IncSubtensor], + ids=lambda op: op.__name__, +) +def test_core_matrix_keys_carry_the_bundled_rules(key, op_type): + """The bundle stays in step with what the built-in matrix properties register. + + ``Subtensor`` is excluded: ``SELECTION`` registers its own rule there instead of the + shared one, so that Op is deliberately not uniform across the built-in keys. + """ + probe = AssumptionKey("probe") + register_matrix_property_rules(probe) + + bundled = set(ASSUMPTION_INFER_REGISTRY[(probe, op_type)]) + assert bundled <= set(ASSUMPTION_INFER_REGISTRY[(key, op_type)]) From 1a4cc21deaad5f3640145106b8ab8d0b8c67daa0 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 22 Aug 2026 23:52:11 -0400 Subject: [PATCH 6/9] Appease mypy --- pytensor/assumptions/core.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/pytensor/assumptions/core.py b/pytensor/assumptions/core.py index 3b0d1ed8c4..9d70d92407 100644 --- a/pytensor/assumptions/core.py +++ b/pytensor/assumptions/core.py @@ -1,5 +1,5 @@ from collections import deque -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Iterator from dataclasses import dataclass from enum import IntFlag, auto from typing import Any @@ -101,7 +101,7 @@ def __reduce__(self): def __repr__(self) -> str: return self.name - def assume(self, x: TensorLike, *, state: bool = True) -> Variable: + def assume(self, x: TensorLike, *, state: bool = True): """Return a view of *x* declaring this assumption. Parameters @@ -208,22 +208,22 @@ def decorator(fn: InferFactFn) -> InferFactFn: return decorator -class KeyRegistryView(Sequence[AssumptionKey]): - """Live, read-only sequence of every registered :class:`AssumptionKey`.""" +class KeyRegistryView: + """Live, read-only view of every registered :class:`AssumptionKey`.""" __slots__ = () - def __getitem__(self, index: int) -> AssumptionKey: - return tuple(KEY_REGISTRY.values())[index] - def __iter__(self) -> Iterator[AssumptionKey]: # Snapshot: a rule that constructs a key would otherwise resize the registry - # mid-iteration. Defining this at all keeps iteration off the inherited - # ``Sequence.__iter__``, which walks ``__getitem__`` index by index. + # mid-iteration. return iter(tuple(KEY_REGISTRY.values())) - def __contains__(self, value) -> bool: - return KEY_REGISTRY.get(getattr(value, "name", None)) == value + def __contains__(self, value: object) -> bool: + # Only a key can be registered, so the isinstance both guards the ``name`` + # access and lets the lookup be a hit rather than a scan. + return ( + isinstance(value, AssumptionKey) and KEY_REGISTRY.get(value.name) == value + ) def __len__(self) -> int: return len(KEY_REGISTRY) From e4d6e7a490c6368a2d18e5cbb6b300e8112f8fa7 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 01:14:39 -0400 Subject: [PATCH 7/9] Document pytensor.assumptions in the library reference --- doc/library/assumptions.rst | 74 +++++++++++++++++++++++++++++++++++++ doc/library/index.rst | 1 + 2 files changed, 75 insertions(+) create mode 100644 doc/library/assumptions.rst diff --git a/doc/library/assumptions.rst b/doc/library/assumptions.rst new file mode 100644 index 0000000000..8ebd73cf7c --- /dev/null +++ b/doc/library/assumptions.rst @@ -0,0 +1,74 @@ +.. _libdoc_assumptions: + +============================================================================== +:mod:`assumptions` -- Structural Assumptions and Assumption-Driven Rewrites +============================================================================== + +.. module:: pytensor.assumptions + :platform: Unix, Windows + :synopsis: Track structural properties of tensors and let rewrites exploit them + +The :mod:`pytensor.assumptions` module records structural facts about symbolic +tensors -- that a matrix is diagonal, triangular, symmetric, positive-definite -- +so that graph rewrites can replace an expensive operation with a cheaper +specialized one without inserting runtime checks. + +Facts are attached to ``(variable, property)`` pairs inside a +:class:`~pytensor.graph.fg.FunctionGraph`, inference is lazy and cached, and an +answer of *unknown* is both common and legitimate. + +For a worked introduction, see :doc:`the assumptions gallery notebook +`. + +Declaring assumptions +===================== + +.. autofunction:: pytensor.assumptions.assume + +.. autoclass:: pytensor.assumptions.SpecifyAssumptions + +Inspecting assumptions +====================== + +.. autofunction:: pytensor.assumptions.check_assumption + +.. autoclass:: pytensor.assumptions.AssumptionFeature + :members: get, check + +.. autoclass:: pytensor.assumptions.FactState + +.. autoclass:: pytensor.assumptions.ConflictingAssumptionsError + +.. autofunction:: pytensor.assumptions.summarize_assumptions + +.. autofunction:: pytensor.assumptions.assumption_tags + +Properties +========== + +Each property is an :class:`AssumptionKey`. The built-in keys are +``DIAGONAL``, ``LOWER_TRIANGULAR``, ``UPPER_TRIANGULAR``, ``SYMMETRIC``, +``POSITIVE_DEFINITE``, ``ORTHOGONAL``, ``SELECTION``, ``PERMUTATION``, and +``UNIQUE_INDICES``. ``MATRIX_KEYS`` holds the eight that describe a matrix; +``ALL_KEYS`` is a live view of every registered key, including those added by +downstream libraries. + +.. autoclass:: pytensor.assumptions.AssumptionKey + :members: assume, holds + +Defining a new property +======================= + +Constructing an :class:`AssumptionKey` registers it, after which +:func:`assume` accepts it by name and ``debugprint(print_assumptions=True)`` +reports it. The functions below say how the new property behaves. + +.. autofunction:: pytensor.assumptions.register_assumption + +.. autofunction:: pytensor.assumptions.register_matrix_property_rules + +.. autofunction:: pytensor.assumptions.register_universal_assumption + +.. autofunction:: pytensor.assumptions.register_implies + +.. autofunction:: pytensor.assumptions.register_constant_inference diff --git a/doc/library/index.rst b/doc/library/index.rst index 6b5dfe29fe..0f73a090d4 100644 --- a/doc/library/index.rst +++ b/doc/library/index.rst @@ -15,6 +15,7 @@ Modules .. toctree:: :maxdepth: 1 + assumptions compile/index config d3viz/index From 3954eafbabc53aee341c897a39147834c038ce8e Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 01:15:30 -0400 Subject: [PATCH 8/9] Add assumptions gallery notebook --- doc/_thumbnails/rewrites/assumptions.png | Bin 0 -> 12626 bytes doc/gallery/rewrites/assumptions.ipynb | 1188 ++++++++++++++++++++++ 2 files changed, 1188 insertions(+) create mode 100644 doc/_thumbnails/rewrites/assumptions.png create mode 100644 doc/gallery/rewrites/assumptions.ipynb diff --git a/doc/_thumbnails/rewrites/assumptions.png b/doc/_thumbnails/rewrites/assumptions.png new file mode 100644 index 0000000000000000000000000000000000000000..f7c35837877ebee0d82c9cf485d211c8ee4ed0b9 GIT binary patch literal 12626 zcmdseXE9!h&GH-q6g7C z4-#!K`e2lI`L{3U<2mp3zSnVGv)OyCz1DrNyZ`R-MpIpdmWq`Mf*{(*j~-}25Rt+8 zhk^_|X?XEz3jC7xR5tL`c75ULWAV}kQnT=Mb8_``vbVhIZS&H@-ql4+_|`4qTY^{Z zJU!h!q(wxW|1(3_^`))I%L(EGunMKyBSQ}eqPIBz!1CmC?IB3C;qim}I=-pv)0DoZ zx@q#8J2&l7u1u;v*0c8;Ef_C6gtPZwPEckbzwD`b+2-`?mupcn(!5cxr1@{hcEogC zjgp~~Qn>$YUrmBf=vQL)ceRup!;Q1M`m>+D*{A%ujtiK@@A#$Pvv--93V1n5z~a+A zHvFC#CfdLtgf;8HI$#?{4VjZHHxWU+s*Oi|XUBcWc8-8BW>ep+fJs|I#lz^?+?|s}qsj<#lhS*!H=!+n1)QNb4AJ>4RmE4vS{4A$~q3m(2rcPyGQ1Cos z1_e21pxboXhcDF`6XFF*VVnnzta)-a;Uqzs3xWDqnd@-ze@rM)A zA@{mmP(zNT4>?fsU0|DKG2MQj6t)m2FE{14lt#=UfcgFY$#(7YEvr6RJdH@u%9ED5 znZ1(D4(>Ro&7CqlHjGhy*(>#Q(O;Qf{@BfTNoMvk-M!65w?-IwSjg#8$jsjx8>t?> zl#%S-2Fq(}KPTaFC*?_K!NAj<40;n26D1cOO~I5xX8E(jz8y#ThLW}YN!PyW*5skV zNV(lCw}ZJz{KWoGvA(dlz~9_{=y6xf$@bA+%ItB@tpCu=42Rs#=f2;S%VSn??3#qx z)BWm5fs&eQMW$#}oqB8BQhI<(1%X4|gyShK`%~JG$;p|-vy1`#T9YmO)W@c(5Z ztyymz8XqxUVY@UN-MXCjHI_<+GsqL~ZaKPN@3$lI3G0O|u3mJG3SXP@=%ao6_U+s& zYT*u%)r{lcJkzbNo$qhVHIZ}D;DYgWEr4TAIecy3pLV}9!Yt=VsTrr-K@&Y4Dovq` z=9}pn8b9Ng_U680JG;-%kd0jV6MNsk|03UJEr0o^OVn$1D3I_c$!x|%;LAbAQC1{J zGM=e*Q{1*Ub#sC9Ose$r?9oVdGdQqC)4k>0K@Ix4^(v!--CTM067>EA?1ZjNuh^jM z+VPsLJoVX*{MjHHJ(Z)j9w6oACo1kKDt@dHax_RL+4Cop)r@WOJ71{tnP%@@Cq0R$ z0qd3eRo%jM)@f;J)029orfJTV!+P%F2Wy{g(^qDK7eCv+VB}GsT9cccm?&Qv%(Y!g z^-#xG$?lJFPXe{{XSeN179TG%tdNlEoe99D2^k=-&T1H{sHn{Ju2;{>VOLkXs?%O} z@fj*9Ded>T23L%3#mS!@ICTqEU(u2wKawH(T8hDXWbRqwr+xZj?=F7Y7`@`ZoHgFb zlac!Nve?bK)q=upXL{LPP1894k&i8OGJ+{GJ}y-b!v(Yxx1A@x7iDK>>&v2+zeJay=JqR9kpMS%%FWnnP$+#_((BO>1meg%ah)YWPxW@Tp=cc|?_1*1P z379T!^`VL?UFVzm^CS0a_cWH+85mD!K66<|g52Ja_Tkdu&XVgUnv@z&W5PjIqh6ubFO?s!cjo^pKSbxkhVDt^u|WGi-l;y3iKKIg9l& zaw=0ky$gSRaPn6bBX45UP&7;Oyy#$aIJMG~Cr?f}cVNI5I6Ia_tsvmZt=j z?NsgAh)YOI4c(zxgTtLZdCYi5oID<=lUO}22-n>{Pl&C=%c`eZ{nDRxPVslHRg=YC zUr1YR`xrf^FiFYOZNA#tZFag{pCqua*UC{Ob?e?~`<10jt(0r$npJapgbmn>_=x8W zaIuUt2=OP2abIW)e|P&=5Pj7Q?Y9WJoCpr&lHt*oEun zou@;yRSlN!qI~}HSn6G`JkiG->7Jq++TSp)5>P8SDgHrdU}?Wg(HkhNgs<}yyzbjV z*l-hzoy=dBm!G0vpRI6u{<`ni7$vLw$D&b+Un7V=W6gBNbS5NKlRco;Uv7GI z$7^S-ogWtY2M;U*k3@wje|h-wkHm}hN?C<6+|t?NYhby`5wKQfv2J^Pb=us)mnq5a z^^;qHbJS)XnApGgN_04ttZz)vE8SWf4=j^vRe8C9x@F{CsT*7k=fucv)GUzg?zp=s z*2m8)Us$?Vzu{HdEmF|HsMse*SaTaOv<)9rml2Xz4;;FyZV*!%*A{bm@ML04Jn-~z zX(`UQy%ut(Ys8#PV#4D>ju#HAvi>6hm;hNQpLWjf^C&KtIIgR7Y;CwbTI#KvcZ7Qw zBAKihqB5*#-^d!sI*z)f`mg36YB9QJ)NP;~OTCZ|}kC{TLd>OLrbT7XBE3FpD1|{vFC9l>uvs>b2!L;8v zX#e=cw7+qFtMsu9Z>Rl_@Tz(qUgHR{XK_Ck+C4b>Oq^VeD+-P$-r33P=r$IFjqQS9 zBY>$`alOKcSj~>m?xw2vTP4u{cQdufT6zn^lF*N=PAR`T-b=)q=pnb&p1g+f9k(AP z_^}|)CEr&0_L5mCvw61JgNF^~oC?KnJ}fLeAgH~;WwFQ?q_L`_`KA~wu4u%?;SyD| z1N%`2mowq5Hk_8|$&a7KY!^9|vm&%bsF;xuRA6GUXbyvz!V#p9rqG9}1u6ejI}TeI;uaVr62Fj1oV_;WH|{=@?lhjI+?s+5HI- zu8OowUrlN6nkPh(TGjAtHLeomEE<>D*6m6EEVR1menVuN-gn0)Jl=qyHPdxLxKW|jFV@ir!nFGH`x(Ag^G;;QX^!TdQOr`jsw7MTd7M?Jm~Zc7(fhN!u!K^%qQwXU zj^v$dL;iqTCfN$89v(ey{9kz~;@hKpj61?#(W)8r?<``D{BqZMB@GxA_Zo^6AO5(< zipnEh!=#3nsfSpFskEn1_*8*6na?&bXF$zd?Eg1qp#bJPtVR!q6fBAUyBD>5ABuka zCpf{WH@=ez`0y`>ZDn=9M8dZZXyYO&Y2Rk?HlSt?v+R0N+fABxk~RV!@Ft(ho!!W? zGu=DDDarPSl0xyHfY%%3^b?j*9(1>$mY)b}#PCFh{-U$E31EeX|K7AcIZ+XQcGs7W4n~ z0%}SpO@iIqeTDj^>f;>d3?r3QmsSIdO-DxxgJTl6x1QHYE4};v-6!mez0u^ws2s;2 zN8V)dv{piOe|du}QyE&RSU=n8wS290=X=qR(_=1nraaeZW)Zrb#U8fGg?#-+yU{`} z@}iQGlFcD)jquJ7Jds-YcRrdxUN;j{#;~%W4*m5@)L)EgNWlC%b~v6G4nZeko$S!n zTL26Bua`n1Il_2f1OPWM7bD47$}xmG3^hPj2&wm3AebYI+}0b7uuytA-(T{0P|ZdX zaF`IONV2iU3o2JGoR`h9V{#mz1Pts~HYb_ok7P&>-0tKe&883hn|Kbq;6wG-H;%x~ zCQ-+cq6?GleZl(`qQS=p7VDg=;Rsz=nJ;ZM1yM{DyU$er5_KWr1W0c!|7x$Ym z0YHxCt#(e%;GL9Nu+h3}NWjc#O-%4Hdlx4FVy4K}4}RLpF`Dj|Cxu(Z$$EeAK&A)n zk8hU61qs&AUn=duoF!J*-3RklyTgo>@snr9Bhd zxLC4xumqTpytnm(aqjz|dUq*S5_8yN(f>vH5Wlcg-LO|Y=X$#6+Ke?jJKA_|0H3^! zRW14&bh03SYISzhu*I&?P)kW!7a^(td+HwM?^%NgDmTKTOV-ZMTLa| zXGb5k)jO+Z{IP4!t^r2-pbWlScet2R5ptLm(g{-f4=4N8m>zZMm2Ybx|NL2<;kWpC zODMS9>||9ReJ+f1FiRp}f6XZ)Xuq?he!I(Qt~q>s+%_XkTwL7dY|rd00ka4s)E9dW zM{NJ@90yVvJn5gAnJHw?P4{>32K9pN8NiK_f(}|Zdoj#$GVb#~OxuCxkyilh&8BXx z{E7FtO_ELZ)XOLPK*BNsoli{P3mKMM*Mz7k#9?m-pB^vCXN>+uEnIVzP142K8-N6K zC_1N&=NY%G9(n`TUz}r3`=Ol2SL;9xS`kVo%{%E*x4zB#62NA*X9UbyimUwZc%|~< zstJB0w7(a=!N|yHhjrNOByi=4S#rq7)+5GAcTj0~j(!F3*Ea&844zco{`{oaHAuh+ zX@1{(=$VSw=-)l-R)qEHK6RZ|rL>meog6|f=+)I%U39bS~lMPTzV?1`g zRtgkZ<@ycV35RElk~6+mrb=vVY!yK0Ef@`lUL=v8Aj4Mk(S{GEt_eReF@ZY!huBlZ zY}I`ZUb|pU_K$7lPgXSW{5rQ{TB$LDrU!s2=isK`6Ra7|D~6@ll53f-RgO+(4e7Ko zRXbjwTj=i`M`{f!Tf=kUVDCQ=j)%5OyE3!M5W2Xc!s9DU06q$effoY$C2N(kr%(HK z$Y@LsQS}G&(E-k9UHT2>O+6zD*v6caJr;fdQDg|2*4GdM@G}9f!9^^3mb5|v=djzl z8DU~Hf(iL$Y(^5jqDCyoIcnl#(r{+viuO_)!|c(HlU4`}#W15z`Il zE#n_o>+Jf!m||w(Mo#b0e}r5wyQO6D0`N#G=xA5eWvXU#uQWtr(%|g{vw=;JJhlekz12}l2RC3>Cs;v7sC&i_! zZZr4k=K*ZMmThwK(}Yl+d9TZnEB^fSl_nlN;z|X?DX0;%qk*Idznw*|5H!G|BLRvK z0dS#6=c@6X0D^Gc+QpksKmBr5d%~r@bnu%;{8e`&wU-Qll`jXCXi4y51e7WxbEsllP$^PBIJjliK+ywiA6_ zeJXYV$CKu%<=f+EI(ULQS1=0zWm2EpYzy1WW<;QqTyK#7aG~j?#h;%C9W%sU{2cq6 zFM;KEj#>)*rdsm%rhO^b{L;C#ulFS9K|nD*0LPscS;-r0TJ2(Tx|IDSiEjD$X_~b7 z*;wkfRo~VX7`K4whajcFA8#0UP|wl@q{pO@`HpDY8vf|YJj^;1FV zoJX7Gr)Z2vDjR(1eq~h^ojvBy&afD7VyiQTFpcwfj?OAC6{;G$Hsv*Hq9iS+Q(@ur z_8<6z{QU2qW7Hmh67(ZR8}=&?(+19=?~S@zpZW#%xuS|*DgZTbw)BOZ?gU!er0$OE zuMMyEwN)iin-v6oBsN%ES)HuY8!@w3)2NqWd;PLZL?fR3p*paI8SFp>#}AG?-v#}K zQ*tA639n-MT#nXvh_TRVGw69kpx>vv9xqs4bDk`#%*{n|(w*DuDMF*~@o~#DIgK)( zB3F}h>v6CyIdUws0&vrG=nB8YTMZN^`VSOWht5OF%R@d3I?;6 z{*fvYACne+Ki^|dBOE7pv`mJc0nl@X1cd=x;p0NR;x75VG@0tr{q#WziR!r1kA0)y zB$GzZ(zzt|O%pKmdv*pM$oY#z?wz5&Ji7a!CGcmutxREWoh)|e%ltx{d%bM0UtU|S za#^N{TNA0Dh!Ks9Ebhtzf17Qp`@_>^&+h=?T{6LK`EYF9Iy;#?lg$|1n!z_z2T)w& zkP%*CSh9oJ|9HN89!|`7rK0v$r*XkPzWvdrA2nQAoY-0H%Us%Q7*97N<$s&CADZDp z&Imh4gZfi&X88@a`1z^a0~ydkG3yGze(`0T^4VC++3(L#qL;lxUyyM$l#7ZawmOND zoG4R4E~!ZjrC5*_`Es>Ii|!k=9_qW(J*%IL_9&dczTaT^EXuL02IGsNP~`Ki#t09&LS-=02TNBc#5%R!cWx-bz{S-k?(dHO6DI3dGT4I%)Wz zlqAU`Zd1&;=yy72+|xBvnVBNzE&KdalkNa949riC{HQH&CQ>hDVT zw0tQ^!I6b@M*{%L^oKqXuD#AF`&!n_{dA4FVKW?tv$85_I07kVQ%$e#nsf9{2Z%^j z$Mq{sU(yrizlQO7Gi9LO_z1fX9JNtttmVWibHQJa-I&z})R}a-03?WX|gUwOpaL>KBxh?i`K+$|{Olydh z1_x0nGc)cjT5F3nCR;>(J9e)C*CG*`f9}1kT!ZQ@R_o5Bn8ka|R%NW2tJ_cTl6_($ zV^Thq#?12!7EG>fM|6L+%?dO~^wMI-2QRt4Ucg&YL~(9KNOFi(aR;MogMRoKZ^6Et zPB5F*ft>glzlsZVSu_4I*K2?)x~Nh77>*>HFYgxa=8S9D9eA>%NUoWWAQgT8usGm) zuYL_^G;GUcH#k0nK|}AKeYW9P5}XR$={aP0o@&dKp#_8TbidwyZ~*#izvuXS*%?X@ zs^dd}e2C^=Q1$1Lb37ydXYHSt_;|xzUDsHNzhP7-Uq^M4Kl5hmWlE6d_l_#E_K45? zJKff!(LX_pNjEP3XuUyxl(WK@_wST_yY3q&Xe3!YwX*UZE*H0mto*t&W8C7t?-+%D z(bdg~5@71_qTZlZ-Z`%@zu4R7pE{Q0j+FFn#$Jl^9*yhQk~1}a6Zk!j_ISLWdzLZk zI@XCk4R3bqck->@2we{N)YnR}d$)s(YrHpxL|yEkjM=MsF(0=rQfAqysz`G8-g&LOy#3e>`(X9K|UvsA$GI|>g!^^e!jKv(a zeC;f6X@U0jaXs#qkxz^h3t98!>Vd{q1O zM6ZNeEo7x3DkD*n*&~<&YIKX~P`q=?cv{p-@u2rOf6*ESdD;HZJ92b(65V2l@}6*C zM5+eeYzROxKzTgeBca3i6-6#;D36)@=!Ll=9MX(?FE}U~AAf0DzV`0-xEVqn5j6Co zarJWf{_|m`#bPYWTO}e$=V^-_N1knsgS}w53J(+RWEFxWa<3747t?0DwNB;Ft&V>v z{y+*1rACnym6;fRsw`<41XQaPJZjwb1hce>}Zkq2g zMN);7?_GCgpoN&UbjgZp9*&e0jvDb$L%iP#XbckFz(Pg*#Cq4kc_`~v!W^x*9LLAn zR3)jQ%%dcGmSt(f9+S^r-T80#B1!l2BTRGi2zE^LSG z$nHc{1CwG!5qMc4E0RKT@$+JvTc^aJtO~Q_z?xYO%07 zgEK0pu)-i3T?uM$rI0`qgW8uipKjZ9<*`>AmnQEAYV(T2mwM*(h3jhu)DWmj(6Kx( zS-D0FlhZ;g`Fwgn?3}Uh-aHs8OK_orG(%T~&07g4305iFc;R}BraM4YAC+tn>7~vD zcg!bP5H1x1ofW+XrX_^fgFzNPugl15|M){3^Rz`lQEzv7cgw}*ztR@v^%sbj`y=t&!LAwa( zGEF*SSVdbOo>S>{>_-{toI_~{HD*#8q!=~d?;wK69yd;OJiY)9H0vUg>_R#G^Z7dp z(;|bSzr4!RlZUS40ORSnOA4)~BE%cxhoy74o;u-J(=eBy%wIGY=r*tK(dfIs6?`HE zL2TUSv;!Us$<3u^ccy{iG(My>2)z?zqT<`%<#n#j5~_`?*Ux9BfSG>*<#LKbB#2rI zWr)7Wt0RUSUfZx=aN3$(uX76#DL!MJY?_lyMsD zSTh!Ah%Zg}jqL59KwrDIFt?Uh_Hb@eg!B#b8&xe)BS5E%59pvhl}6U^ThjH(av%fK4|DV%<ikJa^R^rl58a$&U^T1>B0OG#31N%v=s-dQ%$cUM%03anL$_*bgae8RleBDUST! zOS2fY9eq>>X>#wH>o%D+%5^_)uIa@H$5i@1wNjoWfoS3ox;7kBmVq@6OG4=TcOa%u z2wi1M&XH}(o~*(wCYP}n$zUha8|J$2v`(v}WfR>Pz<%*0MGj`o7Af{~D}*!6l@~B5 z^s4!#tJc$pC#BSekfuiW9l?<05={qIwkXSS-O;cwWBG*wMRi95o4=>4BOF7#J$z4d zhO>A%q%5EDCRdaMiZbS2*)O`UXXz zx~;0Tv=*fyQsPrpHFv+qr0B0KLJ+a_Yb3-~Xmfnbpz2N)NReF>-7i!ASu^=s&$9z&WRNQ(R|L1baMM5KoX_H6!x(qJ!7 zqa(R;e(~2(ZPh?A5hB^*eLz}dk>JPXiLwesm4z_XY6{390-@WdN^3E`8|v@cjoMy^ z8h;&Td~|u6R&0@C8@#%r2(jFRMY{}+9N&srH?Tk7w#)j|68&jKO zVtyUc9vRNfu!`J>Yc-y;-2Z~SAE8F7^+3Cz{N-FL=%?s9tTx$(eK9n&Jxkv+f2U)m z#8*%FJFU5UQJvbvkb-%pjQ4ndC875=^u{XNFp!1WYEspo4#v<`F<`buw?qUt*UENXEb* zhgd{w9#FA7U@Q5LH1zz^ipD;8^;Hy}2tk_KyXMdO(~G1IgCgx2kIx^F2=cR*kNCW> zF#Zs$135&Wjpt3tH)!SS*PxEx{h84tM61sOq8Z+xqHqR&9R}5j3pa_yx&@?WyZosdOGbioNFvQ^>B}%e^^O z8|dYW_o9E!Lyai7YW|Ed7zos0ble8Sm4a6YWi?-Hx(7%{6^0p*d6dAvWF}J>4AB6w zIB=Oi8|1>$ckAf_-iEmPzH3)-IW7s^jA&VMeg<5tW+XVc3NGwxQZYrKO=l>9?&cPO z5j~V4G(-uDPG)Wy9^Ej#*o*%2&J8;#q60G2As8$g6>noG$Af+1bC1g!COZGK>5I#v zwX&q#p(1db#I_oRZrz|JB7%E?%)$H)rNIy_Ag2GC`Lk*+$2o0*mQQ4Z46sKUjbB`p zMhdH@Z#zkPDQ2e-E6~0qs82@}rO!F6Bv*-|?bSdeQ& z5>fKtdyq{Ga|@F9IBwp#6Y=F>HFya$96(+Rr8)l zjxUM<0b@j(i!GMe<3R}#dd<8E621s>AW;?~IL)Raup<^;7-pL0rQGdZf zNXeMT(IG&8%rYod z0vYI@QsXPC@FMcux76tZBzbCy_6Vx5qR&|AW)6Ypp*D|T4j(uK^BR@$^GilbgJ!7G z`CQvJ)F>?P1`%9GlSW_CS(Y!*wQdC*_Yeb&n*R#&Iop$TVUoOCVElk+*ON-Y+TUO0_SZV(0>yU(E<13Iey8mGqUYc%w>)PIK@4m)prKx zITYGQp{)I5=O23hmuyuDcc7!6BkvzGHQomuM0w;lYnA>&+c{Z{g$#W9bxEar)$v`Qe>D0Y&YYqa$7XWVD6yXYc;QSscjW9ZcmOor!`gSL~gP^tY0;{S; za0P(}p$dk{l=K1%kUB>@72$@7UT2pP{c~_;2T|(;VqKNQa94j6gZa=((<@f*!GgWH zeR1}=UhyqwFQ6l$&w)A8%6O_)0DsD_(D;h=1tmG6ivgM029)EgM&m1K8p9;nI*B5o zzu7W~nZh4$XxefW*6L7&h}i>J+iryn@$Q=W{j~6+52xiU_J&6K@HPnwBZ%+vwJ_x z1QceMM9@4melJDWnXCJpR-Ht9$UnSjyoYw~k9J#xR_6dMq?Aogec*VXH=((Iur1tO zcBJ(=s$uXl5zBQL^;lj;u>rGO*|%1Zp@d#Hhn+-uAQLNn{&| zZ)~JAXl$bnizgMmP6S7;Tmx?R+~FW#kwweZLhBYbwpZ?um|FpYG`Sql9$PiLFt$lK zLMi628CY~WRxg)~c)plBtz7TDs=!GE7rkC=Vpy}elIXTIPzfw^S*ejI!rF;1>~7Cn z5I&Y85u`M4Kva`Z2hX2&5G`jQNUzfy`V`;VV$X90T)r>)uxIY5@=3f6&89g^B6!ul zG6{Q72f2kN}((})Ce<9e&F9taF!OA(u@hFi3yqJVyv|P@Hy@^vpQU; zvRA{TG*T81!GxyC9,\n", + " 'symmetric': ,\n", + " 'positive_definite': ,\n", + " 'orthogonal': }" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "facts(x_diag, DIAGONAL, SYMMETRIC, POSITIVE_DEFINITE, ORTHOGONAL)" + ] + }, + { + "cell_type": "markdown", + "id": "0ec6f1cd", + "metadata": {}, + "source": [ + "The distinction between `UNKNOWN` and `FALSE` matters here. The system cannot determine whether a diagonal matrix is positive-definite (a diagonal matrix may or may not be), so it reports `UNKNOWN`, not `FALSE`. `.check()` collapses both to Python `False`, so use `.get()` when the distinction is relevant.\n", + "\n", + "We asserted only `diagonal=True`, yet `symmetric` is reported `TRUE`. That follows from the implication system, described next." + ] + }, + { + "cell_type": "markdown", + "id": "d36eae88", + "metadata": {}, + "source": [ + "## Implications: assert the strongest fact\n", + "\n", + "Properties are connected by a small, explicit implication lattice. A diagonal matrix is symmetric and both-triangular; a positive-definite matrix is symmetric; a permutation matrix is orthogonal and a selection matrix. The `IMPLIES` registry holds these edges:" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "01adcd26", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "diagonal => lower_triangular, upper_triangular, symmetric\n", + "positive_definite => symmetric\n", + "permutation => selection, orthogonal\n" + ] + } + ], + "source": [ + "for stronger, weaker in IMPLIES.items():\n", + " print(f\"{stronger.name:18s} => {', '.join(w.name for w in weaker)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "bc2cb969", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'diagonal': ,\n", + " 'symmetric': ,\n", + " 'lower_triangular': ,\n", + " 'upper_triangular': }" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# A single diagonal=True assertion answers all four of these as TRUE.\n", + "facts(x_diag, DIAGONAL, SYMMETRIC, LOWER_TRIANGULAR, UPPER_TRIANGULAR)" + ] + }, + { + "cell_type": "markdown", + "id": "8fc5c6d2", + "metadata": {}, + "source": [ + "Implication runs in both directions: forward (a stronger `TRUE` makes the weaker facts `TRUE`) and contrapositive (a weaker `FALSE` makes the stronger fact `FALSE`; a matrix that is not symmetric cannot be diagonal). In practice, assert only the strongest property you know and let the weaker ones follow." + ] + }, + { + "cell_type": "markdown", + "id": "933eb881", + "metadata": {}, + "source": [ + "## Conflicts are caught\n", + "\n", + "Because `assume(..., property=False)` records genuine `FALSE` evidence, asserting something the system can *prove* wrong produces a `CONFLICT`, raised as a {class}`~pytensor.assumptions.ConflictingAssumptionsError` the moment the fact is queried. For example, `pt.eye(5)` is provably diagonal, so asserting that it is *not*:" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "cb24e78e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ConflictingAssumptionsError: Conflicting evidence for diagonal on SpecifyAssumptions{!diagonal}.0 from owner-inferred rules.\n" + ] + } + ], + "source": [ + "contradiction = assume(pt.eye(5), diagonal=False)\n", + "\n", + "try:\n", + " facts(contradiction, DIAGONAL)\n", + "except ConflictingAssumptionsError as err:\n", + " print(\"ConflictingAssumptionsError:\", err)" + ] + }, + { + "cell_type": "markdown", + "id": "4dddbd6a", + "metadata": {}, + "source": [ + "## Facts propagate through the graph\n", + "\n", + "You annotate the inputs, not every intermediate result, and the system carries properties forward. Inference walks the graph inputs-first, and each operation has rules describing what it does to a property: the {func}`~pytensor.tensor.linalg.cholesky` factor of a diagonal matrix is diagonal, a transpose swaps lower- and upper-triangular, and the product of two diagonals is diagonal. Constants are inspected directly, so a literal identity matrix is recognized as diagonal, orthogonal, and a permutation.\n", + "\n", + "Here the `diagonal` fact flows from `A` through `cholesky` to `L` with no annotation on `L` itself:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "0bf627ff", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Blockwise{Cholesky{lower=True, overwrite_a=False}, (m,m)->(m,m)} [id A] a={diag}\n", + " └─ SpecifyAssumptions{diagonal} [id B] a={diag}\n", + " └─ A [id C]\n" + ] + } + ], + "source": [ + "A = assume(pt.matrix(\"A\", shape=(3, 3)), diagonal=True)\n", + "L = pt.linalg.cholesky(A)\n", + "\n", + "dprint(L, print_assumptions=True);" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "d566e835", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "U {'lower_triangular': , 'upper_triangular': }\n", + "U.T {'lower_triangular': , 'upper_triangular': }\n" + ] + } + ], + "source": [ + "# Transposing swaps lower- and upper-triangular.\n", + "U = assume(pt.matrix(\"U\", shape=(3, 3)), upper_triangular=True)\n", + "\n", + "print(f\"{'U':<5}{facts(U, LOWER_TRIANGULAR, UPPER_TRIANGULAR)}\")\n", + "print(f\"{'U.T':<5}{facts(U.T, LOWER_TRIANGULAR, UPPER_TRIANGULAR)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "7bd2a612", + "metadata": {}, + "source": [ + "## Rewrites that consume assumptions\n", + "\n", + "The preceding machinery exists to enable rewrites. Compiling a function with `mode=\"FAST_RUN\"` runs rewrites that query the fact cache and specialize the graph; the `SpecifyAssumptions` markers are removed in the process. Each of the four examples below produces a compiled graph in which an expensive operation has been eliminated." + ] + }, + { + "cell_type": "markdown", + "id": "f33e0d68", + "metadata": {}, + "source": [ + "### Diagonal matmul becomes an elementwise product\n", + "\n", + "A general matrix multiply is $O(n^3)$. If both operands are diagonal, the product is diagonal and reduces to the elementwise product of the two diagonals. No `Matmul` remains in the compiled graph." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "3df8c988", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "FusedElemwise{Mul} [id A] 3\n", + " ├─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id B] 2\n", + " │ └─ d1 [id C]\n", + " ├─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id D] 1\n", + " │ └─ d2 [id E]\n", + " ├─ [0 1 2] [id F]\n", + " ├─ [0 1 2] [id F]\n", + " └─ Alloc [id G] 0\n", + " ├─ 0.0 [id H]\n", + " ├─ 3 [id I]\n", + " └─ 3 [id I]\n", + "\n", + "Inner graphs:\n", + "\n", + "FusedElemwise{Mul} [id A]\n", + " ← AdvancedSetSubtensor [id J]\n", + " ├─ i4 [id K]\n", + " ├─ Mul [id L]\n", + " │ ├─ i0 [id M]\n", + " │ └─ i1 [id N]\n", + " ├─ i3 [id O]\n", + " └─ i3 [id O]\n" + ] + } + ], + "source": [ + "d1 = pt.matrix(\"d1\", shape=(3, 3))\n", + "d2 = pt.matrix(\"d2\", shape=(3, 3))\n", + "product = assume(d1, diagonal=True) @ assume(d2, diagonal=True)\n", + "\n", + "f_product = function([d1, d2], product, mode=\"FAST_RUN\")\n", + "dprint(f_product);" + ] + }, + { + "cell_type": "markdown", + "id": "3aa5b33c", + "metadata": {}, + "source": [ + "### Orthogonal $Q Q^\\top \\to I$\n", + "\n", + "For an orthogonal matrix, $Q Q^\\top$ is the identity. The product reduces to a constant, and the multiply is removed entirely." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "9b8d2584", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DeepCopyOp [id A] 0\n", + " └─ [[1. 0. 0. ... 0. 0. 1.]] [id B]\n" + ] + } + ], + "source": [ + "q = pt.matrix(\"q\", shape=(3, 3))\n", + "q_orth = assume(q, orthogonal=True)\n", + "gram = q_orth @ q_orth.T\n", + "\n", + "f_gram = function([q], gram, mode=\"FAST_RUN\")\n", + "dprint(f_gram);" + ] + }, + { + "cell_type": "markdown", + "id": "2d4adfc0", + "metadata": {}, + "source": [ + "### Positive-definite solve becomes a Cholesky solve\n", + "\n", + "{func}`~pytensor.tensor.linalg.solve` for a general `A` uses an LU-based solver. If `A` is positive-definite, a Cholesky-based solver is roughly twice as fast and more numerically stable. The assumption selects the specialized path, visible as `assume_a='pos'` on the compiled `Solve`." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "aa1be802", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Solve{assume_a='pos', lower=False, b_ndim=1, overwrite_a=False, overwrite_b=False} [id A] 0\n", + " ├─ A_pd [id B]\n", + " └─ b [id C]\n" + ] + } + ], + "source": [ + "A_pd = pt.matrix(\"A_pd\", shape=(3, 3))\n", + "b = pt.vector(\"b\", shape=(3,))\n", + "solution = pt.linalg.solve(assume(A_pd, positive_definite=True), b)\n", + "\n", + "f_solve = function([A_pd, b], solution, mode=\"FAST_RUN\")\n", + "dprint(f_solve);" + ] + }, + { + "cell_type": "markdown", + "id": "a1fc4cb4", + "metadata": {}, + "source": [ + "### Kronecker product of diagonals\n", + "\n", + "The {func}`~pytensor.tensor.linalg.kron` product of two diagonal matrices is itself diagonal, so the dense `KroneckerProduct` op is replaced by a construction from the outer product of the two diagonals. No `KroneckerProduct` remains in the compiled graph." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "093560fd", + "metadata": { + "scrolled": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "AdvancedSetSubtensor [id A] 7\n", + " ├─ Alloc [id B] 6\n", + " │ ├─ 0.0 [id C]\n", + " │ ├─ 12 [id D]\n", + " │ └─ 12 [id D]\n", + " ├─ Reshape{1} [id E] 5\n", + " │ ├─ Mul [id F] 4\n", + " │ │ ├─ ExpandDims{axis=1} [id G] 3\n", + " │ │ │ └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id H] 2\n", + " │ │ │ └─ k1 [id I]\n", + " │ │ └─ ExpandDims{axis=0} [id J] 1\n", + " │ │ └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id K] 0\n", + " │ │ └─ k2 [id L]\n", + " │ └─ [-1] [id M]\n", + " ├─ [ 0 1 2 ... 9 10 11] [id N]\n", + " └─ [ 0 1 2 ... 9 10 11] [id N]\n" + ] + } + ], + "source": [ + "k1 = pt.matrix(\"k1\", shape=(3, 3))\n", + "k2 = pt.matrix(\"k2\", shape=(4, 4))\n", + "kron = pt.linalg.kron(assume(k1, diagonal=True), assume(k2, diagonal=True))\n", + "\n", + "f_kron = function([k1, k2], kron, mode=\"FAST_RUN\")\n", + "dprint(f_kron);" + ] + }, + { + "cell_type": "markdown", + "id": "e2f1232f", + "metadata": {}, + "source": [ + "## Worked example: a Gaussian process marginal likelihood\n", + "\n", + "Gaussian process regression is a natural setting for structural assumptions, because its central object is a covariance matrix that is symmetric and positive-definite by construction.\n", + "\n", + "For inputs $X$, targets $y$, a kernel $k$, and observation noise $\\sigma^2$, the training covariance is\n", + "\n", + "$$K = k(X, X) + \\sigma^2 I,$$\n", + "\n", + "which is symmetric positive-definite. The log marginal likelihood is\n", + "\n", + "$$\\log p(y \\mid X) = -\\tfrac{1}{2}\\left(y^\\top K^{-1} y + \\log|K| + n \\log 2\\pi\\right),$$\n", + "\n", + "and a direct translation writes $K^{-1}$ as `inv(K)` and $\\log|K|$ as `slogdet(K)`. This is the pattern GP libraries built on PyTensor use: a kernel annotates every covariance $k(X, X)$ as symmetric and positive-definite when it is constructed, the likelihood and posterior code is written in this naive form, and the assumption system specializes it." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "765aeb80", + "metadata": {}, + "outputs": [], + "source": [ + "def linalg_op_counts(fn):\n", + " \"\"\"Count the linear-algebra ops remaining in a compiled function.\"\"\"\n", + " keep = (\"Cholesky\", \"Solve\", \"MatrixInverse\", \"SLogDet\", \"Det\", \"LU\")\n", + " counts = {}\n", + " for node in fn.maker.fgraph.apply_nodes:\n", + " name = type(node.op).__name__\n", + " if any(k in name for k in keep):\n", + " counts[name] = counts.get(name, 0) + 1\n", + " return counts" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "7f9ecb20", + "metadata": {}, + "outputs": [], + "source": [ + "def exp_quad_cov(X, ls):\n", + " \"\"\"Squared-exponential (RBF) covariance k(X, X).\"\"\"\n", + " sq_dist = ((X[:, None, :] - X[None, :, :]) ** 2).sum(-1)\n", + " return pt.exp(-0.5 * sq_dist / ls**2)\n", + "\n", + "\n", + "X = pt.matrix(\"X\")\n", + "y = pt.vector(\"y\")\n", + "ls = pt.scalar(\"ls\")\n", + "sigma = pt.scalar(\"sigma\")\n", + "n = X.shape[0]\n", + "\n", + "K = exp_quad_cov(X, ls) + sigma**2 * pt.eye(n)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "66b3d5bb", + "metadata": {}, + "outputs": [], + "source": [ + "def marginal_log_likelihood(cov):\n", + " quad = y @ pt.linalg.inv(cov) @ y # y^T K^-1 y\n", + " _, logdet = pt.linalg.slogdet(cov) # log|K|\n", + " return -0.5 * (quad + logdet + n * np.log(2 * np.pi))" + ] + }, + { + "cell_type": "markdown", + "id": "b9ad04c7", + "metadata": {}, + "source": [ + "Compiled from the plain covariance `K`, the linear algebra falls back to general routines. PyTensor already avoids forming an explicit inverse, but nothing tells it that `K` is symmetric or positive-definite, so the solve and the log-determinant become two independent general (LU-based) factorizations:" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "76f57f51", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'SLogDet': 1, 'Solve': 1}" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "f_plain = function([X, y, ls, sigma], marginal_log_likelihood(K), mode=\"FAST_RUN\")\n", + "linalg_op_counts(f_plain)" + ] + }, + { + "cell_type": "markdown", + "id": "e455be2c", + "metadata": {}, + "source": [ + "Now assert what a GP library knows at construction time: the covariance is symmetric and positive-definite. This is a single call, exactly the annotation a kernel attaches to `k(X, X)`:" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "bd0f9834", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'CholeskySolve': 1, 'Cholesky': 1}" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "K_pd = assume(K, positive_definite=True, symmetric=True)\n", + "\n", + "f_pd = function([X, y, ls, sigma], marginal_log_likelihood(K_pd), mode=\"FAST_RUN\")\n", + "linalg_op_counts(f_pd)" + ] + }, + { + "cell_type": "markdown", + "id": "67cc6e6f", + "metadata": {}, + "source": [ + "The two general factorizations collapse to a single `Cholesky`. Both the inverse-solve (`CholeskySolve`) and the log-determinant reuse that one factor $L$: since $K = L L^\\top$, the term $K^{-1} y$ is obtained by triangular solves and $\\log|K| = 2\\sum_i \\log L_{ii}$. A Cholesky factorization costs about half an LU factorization and is better conditioned for positive-definite matrices, so the assumption improves both speed and numerical stability." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "09845dc5", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(array(-13.35336976), array(-13.35336976))" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rng = np.random.default_rng(0)\n", + "X_val = rng.normal(size=(6, 1))\n", + "y_val = rng.normal(size=6)\n", + "\n", + "# Same value, different graph.\n", + "f_plain(X_val, y_val, 1.0, 0.5), f_pd(X_val, y_val, 1.0, 0.5)" + ] + }, + { + "cell_type": "markdown", + "id": "fc504932", + "metadata": {}, + "source": [ + "### Fitting requires the gradient\n", + "\n", + "Fitting a GP means optimizing the kernel hyperparameters, which needs the gradient of the marginal likelihood. Taking that gradient and inspecting the compiled graph shows the specialization is not yet complete:" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "c63a77c4", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'CholeskySolve': 3, 'MatrixInverse': 1, 'Cholesky': 1}" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "loss = marginal_log_likelihood(K_pd)\n", + "grad_ls, grad_sigma = pt.grad(loss, [ls, sigma])\n", + "\n", + "f_grad = function([X, y, ls, sigma], [loss, grad_ls, grad_sigma], mode=\"FAST_RUN\")\n", + "linalg_op_counts(f_grad)" + ] + }, + { + "cell_type": "markdown", + "id": "f69cd005", + "metadata": {}, + "source": [ + "A `MatrixInverse` survives. Differentiating $\\log|K|$ produces a standalone $K^{-1}$, and PyTensor turns an inverse into a solve only when it sits next to a matmul; a bare inverse of a matrix it cannot otherwise tell is positive-definite is left as a general `MatrixInverse`.\n", + "\n", + "The assumption is still available ({func}`~pytensor.assumptions.check_assumption` returns `True` for that matrix), so a small rewrite closes the gap. When the inverse is applied to a positive-definite matrix, factor it once and solve with {func}`~pytensor.tensor.linalg.cho_solve`:\n", + "\n", + "$$A^{-1} = (L L^\\top)^{-1}, \\qquad L = \\operatorname{chol}(A).$$\n", + "\n", + "{func}`~pytensor.tensor.linalg.inv` batches its input, so the op to match is `Blockwise(MatrixInverse)`, spelled `blockwise_of(MatrixInverse)`:\n", + "\n", + "The next cell is not idempotent: `register_specialize` adds to a global rewrite database, so `f_grad` above is the \"before\" graph only because it was compiled first. Re-running that earlier cell now would compile it with the rewrite applied." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "8b0f7040", + "metadata": {}, + "outputs": [], + "source": [ + "from pytensor.graph.rewriting.basic import node_rewriter\n", + "from pytensor.tensor.rewriting.basic import register_specialize\n", + "from pytensor.tensor.rewriting.blockwise import blockwise_of\n", + "from pytensor.tensor.linalg import MatrixInverse, cholesky, cho_solve\n", + "from pytensor.assumptions import check_assumption\n", + "\n", + "\n", + "@register_specialize\n", + "@node_rewriter([blockwise_of(MatrixInverse)])\n", + "def inv_of_psd_to_cho_solve(fgraph, node):\n", + " \"\"\"Replace inv(A) with a Cholesky solve when A is known positive-definite.\"\"\"\n", + " [A] = node.inputs\n", + " if not check_assumption(fgraph, A, POSITIVE_DEFINITE):\n", + " return None\n", + " L = cholesky(A, lower=True)\n", + " identity = pt.eye(A.shape[-1], dtype=A.dtype)\n", + " return [cho_solve((L, True), identity)]" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "fc91e38b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'CholeskySolve': 4, 'Cholesky': 1}" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "f_grad_fast = function([X, y, ls, sigma], [loss, grad_ls, grad_sigma], mode=\"FAST_RUN\")\n", + "\n", + "assert np.allclose(\n", + " f_grad(X_val, y_val, 1.0, 0.5),\n", + " f_grad_fast(X_val, y_val, 1.0, 0.5),\n", + ")\n", + "\n", + "linalg_op_counts(f_grad_fast)" + ] + }, + { + "cell_type": "markdown", + "id": "2ba0e1b5", + "metadata": {}, + "source": [ + "The `MatrixInverse` is gone, and the count still shows a single `Cholesky`: the factor introduced by the rewrite merges with the one already built for the forward pass, so the whole loss-and-gradient graph shares one factorization. The gradient values are unchanged.\n", + "\n", + "This rewrite is exactly what a GP library ships. In `ptgp`, for example, it is registered once as `matrix_inverse_specialize`, alongside companions that lower $\\det(L L^\\top)$ and $\\operatorname{diag}(A A^\\top)$, so a full marginal-likelihood-and-gradient graph compiles to a single Cholesky and no explicit inverse." + ] + }, + { + "cell_type": "markdown", + "id": "dfe3f0fc", + "metadata": {}, + "source": [ + "## Extending the system\n", + "\n", + "Constructing an {class}`~pytensor.assumptions.AssumptionKey` registers it. From that point the property is a first-class citizen: {func}`~pytensor.assumptions.assume` accepts it by name, `dprint(..., print_assumptions=True)` reports it, and the rewrite that drains declarations into the fact cache resolves it. Nothing else has to be wired up.\n", + "\n", + "What remains is to say how the property behaves:\n", + "\n", + "- {func}`~pytensor.assumptions.register_assumption`: a decorator registering a per-operation inference rule. A rule receives `(key, op, feature, fgraph, node, input_states)` and returns one {class}`~pytensor.assumptions.FactState` per output. Pass `prepend=True` to run ahead of rules already registered for the same pair.\n", + "- {func}`~pytensor.assumptions.register_matrix_property_rules`: install the standard rules for a property of the trailing two axes, so a new matrix property survives transposes, reshapes, indexing and broadcasting without writing any of them.\n", + "- {func}`~pytensor.assumptions.register_implies`: add edges to the implication lattice.\n", + "- {func}`~pytensor.assumptions.register_constant_inference`: infer a fact from the data of a literal constant.\n", + "\n", + "The example below defines an `INVERTIBLE` property, registers that any identity matrix (`Eye`) is invertible, and records that positive-definiteness implies invertibility." + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "3838b4f3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "eye {'invertible': }\n", + "pos_def {'invertible': }\n" + ] + } + ], + "source": [ + "from pytensor.assumptions import AssumptionKey, register_assumption, register_implies\n", + "from pytensor.tensor.basic import Eye\n", + "\n", + "INVERTIBLE = AssumptionKey(\"invertible\", \"inv\")\n", + "\n", + "\n", + "@register_assumption(INVERTIBLE, Eye)\n", + "def _eye_is_invertible(key, op, feature, fgraph, node, input_states):\n", + " return [FactState.TRUE]\n", + "\n", + "\n", + "# A positive-definite matrix is always invertible.\n", + "register_implies(POSITIVE_DEFINITE, INVERTIBLE)\n", + "\n", + "pos_def = assume(pt.matrix(\"m\", shape=(3, 3)), positive_definite=True)\n", + "\n", + "# ``eye`` is invertible by the rule above; ``pos_def`` follows from the implication.\n", + "print(f\"{'eye':<10}{facts(pt.eye(4), INVERTIBLE)}\")\n", + "print(f\"{'pos_def':<10}{facts(pos_def, INVERTIBLE)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "57e0ae55", + "metadata": {}, + "source": [ + "Because the key registered itself, `assume()` takes `invertible=True` alongside the built-in keywords, and the key can declare and query the fact directly:" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "aae66c4f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "assume(m, invertible=True) True\n", + "INVERTIBLE.assume(m) True\n", + "m False\n", + "SpecifyAssumptions{invertible} [id A] a={inv}\n", + " └─ m2 [id B]\n" + ] + } + ], + "source": [ + "m = pt.matrix(\"m2\", shape=(3, 3))\n", + "\n", + "# The first two spell the same declaration; the third is the undeclared control.\n", + "spellings = {\n", + " \"assume(m, invertible=True)\": assume(m, invertible=True),\n", + " \"INVERTIBLE.assume(m)\": INVERTIBLE.assume(m),\n", + " \"m\": m,\n", + "}\n", + "\n", + "for label, var in spellings.items():\n", + " print(f\"{label:<30}{INVERTIBLE.holds(var)}\")\n", + "\n", + "dprint(INVERTIBLE.assume(m), print_assumptions=True);" + ] + }, + { + "cell_type": "markdown", + "id": "d3bc3ba3", + "metadata": {}, + "source": [ + "### A matrix property: stochastic matrices\n", + "\n", + "The rows of a row-stochastic matrix (a Markov transition matrix) sum to one, and the product of two such matrices is again stochastic.\n", + "\n", + "`register_matrix_property_rules` supplies the plumbing (the property survives batch indexing, reshapes, and broadcasting), leaving only the rule specific to this property:" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "cf5152c9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "P True\n", + "P @ Q True\n" + ] + } + ], + "source": [ + "from pytensor.assumptions import register_matrix_property_rules\n", + "from pytensor.tensor.math import Dot, Sum\n", + "\n", + "STOCHASTIC = AssumptionKey(\"stochastic\", \"stoch\")\n", + "register_matrix_property_rules(STOCHASTIC)\n", + "\n", + "\n", + "@register_assumption(STOCHASTIC, Dot)\n", + "def _product_of_stochastic(key, op, feature, fgraph, node, input_states):\n", + " \"\"\"A product of row-stochastic matrices is row-stochastic.\"\"\"\n", + " if all(state is FactState.TRUE for state in input_states):\n", + " return [FactState.TRUE]\n", + " return [FactState.UNKNOWN]\n", + "\n", + "\n", + "P = pt.matrix(\"P\", shape=(3, 3))\n", + "Q = pt.matrix(\"Q\", shape=(3, 3))\n", + "two_steps = STOCHASTIC.assume(P) @ STOCHASTIC.assume(Q)\n", + "\n", + "# ``@`` is Blockwise(Dot); the delegate every key gets forwards to the core op.\n", + "print(f\"{'P':<10}{STOCHASTIC.holds(STOCHASTIC.assume(P))}\")\n", + "print(f\"{'P @ Q':<10}{STOCHASTIC.holds(two_steps)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "27272598", + "metadata": {}, + "source": [ + "The rewrite reads the fact and replaces the reduction with a constant. Because the fact reached the product, the whole two-step chain collapses: the compiled graph contains neither the sum nor the matmul." + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "b9839f9f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Alloc [id A] 0\n", + " ├─ 1.0 [id B]\n", + " └─ 3 [id C]\n" + ] + } + ], + "source": [ + "@register_specialize\n", + "@node_rewriter([Sum])\n", + "def _stochastic_rows_sum_to_one(fgraph, node):\n", + " \"\"\"Replace a row-sum with ones when the matrix is known stochastic.\"\"\"\n", + " [mat] = node.inputs\n", + " if mat.type.ndim < 2 or node.op.axis != (mat.type.ndim - 1,):\n", + " return None\n", + " if not check_assumption(fgraph, mat, STOCHASTIC):\n", + " return None\n", + " return [pt.ones(mat.shape[:-1], dtype=node.outputs[0].dtype)]\n", + "\n", + "\n", + "f_rows = function([P, Q], two_steps.sum(axis=-1), mode=\"FAST_RUN\")\n", + "dprint(f_rows);" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "3c6edb78", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "True" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "stream = np.random.default_rng(1)\n", + "rows = stream.dirichlet(np.ones(3), size=3)\n", + "cols = stream.dirichlet(np.ones(3), size=3)\n", + "\n", + "# The constant the rewrite inserted is the value the sum would have computed.\n", + "np.allclose(f_rows(rows, cols), (rows @ cols).sum(axis=-1))" + ] + }, + { + "cell_type": "markdown", + "id": "97f87dd3", + "metadata": {}, + "source": [ + "### A property that is not about matrices: sorted vectors\n", + "\n", + "Nothing about the system assumes a property describes a matrix. A vector known to be nondecreasing (a knot vector, bin edges, a grid of quantile levels) makes `sort` redundant, and a contiguous slice of it is still sorted.\n", + "\n", + "The two rules below are the whole definition." + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "dffe0d7e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DeepCopyOp [id A] 1\n", + " └─ Subtensor{start:} [id B] 0\n", + " ├─ grid [id C]\n", + " └─ 1 [id D]\n" + ] + } + ], + "source": [ + "from pytensor.tensor.sort import SortOp\n", + "from pytensor.tensor.subtensor import Subtensor\n", + "\n", + "SORTED = AssumptionKey(\"sorted\", \"sort\")\n", + "\n", + "\n", + "@register_assumption(SORTED, Subtensor)\n", + "def _slice_of_sorted_is_sorted(key, op, feature, fgraph, node, input_states):\n", + " \"\"\"A contiguous slice of a sorted vector is still sorted.\"\"\"\n", + " if input_states[0] is not FactState.TRUE:\n", + " return [FactState.UNKNOWN]\n", + " if all(isinstance(index, slice) for index in op.idx_list):\n", + " return [FactState.TRUE]\n", + " return [FactState.UNKNOWN]\n", + "\n", + "\n", + "@register_specialize\n", + "@node_rewriter([SortOp])\n", + "def _sort_of_sorted_is_a_noop(fgraph, node):\n", + " \"\"\"Sorting an already-sorted vector returns it unchanged.\"\"\"\n", + " if not check_assumption(fgraph, node.inputs[0], SORTED):\n", + " return None\n", + " return [node.inputs[0]]\n", + "\n", + "\n", + "raw = pt.vector(\"grid\", shape=(5,))\n", + "grid = SORTED.assume(raw)\n", + "\n", + "# The slice carries the fact, so the sort is dropped.\n", + "f_sort = function([raw], pt.sort(grid[1:]), mode=\"FAST_RUN\")\n", + "dprint(f_sort);" + ] + }, + { + "cell_type": "markdown", + "id": "3434a4f4", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "- Declare structural facts with {func}`~pytensor.assumptions.assume`. The result is a runtime no-op view of `x`.\n", + "- Facts are three-valued (`TRUE`, `FALSE`, `UNKNOWN`), stored per graph, inferred lazily, and linked by a small implication lattice, so asserting the strongest property implies the weaker ones.\n", + "- Contradictions raise {class}`~pytensor.assumptions.ConflictingAssumptionsError`.\n", + "- Facts propagate through the graph automatically; inspect them with `dprint(..., print_assumptions=True)`.\n", + "- Compiling with rewrites turns those facts into faster graphs: diagonal matmuls become elementwise products, positive-definite solves become Cholesky solves, orthogonal Gram products fold to the identity, and Kronecker products of diagonals collapse.\n", + "- A new property is a custom {class}`~pytensor.assumptions.AssumptionKey`, and constructing one registers it. {func}`~pytensor.assumptions.register_assumption`, {func}`~pytensor.assumptions.register_matrix_property_rules`, {func}`~pytensor.assumptions.register_implies`, and {func}`~pytensor.assumptions.register_constant_inference` say how it behaves, and a `node_rewriter` that calls {func}`~pytensor.assumptions.check_assumption` turns it into a faster graph. A property need not describe a matrix.\n", + "\n", + "The built-in properties are `diagonal`, `lower_triangular`, `upper_triangular`, `symmetric`, `positive_definite`, `orthogonal`, `selection`, `permutation`, and `unique_indices`.\n", + "\n", + "The full API is documented in {ref}`libdoc_assumptions`." + ] + }, + { + "cell_type": "markdown", + "id": "9b256bc6", + "metadata": {}, + "source": [ + "## Authors\n", + "\n", + "- Authored by Jesse Grabowski in August 2026" + ] + }, + { + "cell_type": "markdown", + "id": "4179c367", + "metadata": {}, + "source": [ + "## Watermark " + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "id": "5c81029f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Last updated: Sun, 23 Aug 2026\n", + "\n", + "Python implementation: CPython\n", + "Python version : 3.14.7\n", + "IPython version : 9.16.1\n", + "\n", + "pytensor: 3.3.0+16.g1a4cc21de\n", + "\n", + "numpy : 2.5.2\n", + "pytensor: 3.3.0+16.g1a4cc21de\n", + "\n", + "Watermark: 2.6.0\n", + "\n" + ] + } + ], + "source": [ + "%load_ext watermark\n", + "%watermark -n -u -v -iv -w -p pytensor" + ] + }, + { + "cell_type": "markdown", + "id": "0e5e052b", + "metadata": {}, + "source": [ + ":::{include} ../page_footer.md \n", + ":::" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From acd58ce889b354ce68077e70d0a835e8f81de3ad Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 19:54:41 -0400 Subject: [PATCH 9/9] Add files generated by doc build scripts to gitignore --- .gitignore | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.gitignore b/.gitignore index d2c130f80a..4c49665315 100644 --- a/.gitignore +++ b/.gitignore @@ -65,6 +65,11 @@ coverage.xml # gallery notebook downloaded data doc/gallery/**/data/ +# written on every build by the generate_gallery sphinx extension; a custom +# thumbnail that overrides extraction is committed with git add -f +doc/gallery/gallery.rst +doc/_thumbnails/ + # JupyterLab session artifacts .jupyter/ .jupyter_ystore.db \ No newline at end of file