Make the assumptions system extensible by downstream libraries

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
+</gallery/rewrites/assumptions>`.
+
+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

doc/library/index.rst

@@ -15,6 +15,7 @@ Modules
 .. toctree::
    :maxdepth: 1
 
+   assumptions
    compile/index
    config
    d3viz/index

pytensor/assumptions/__init__.py

@@ -14,10 +14,12 @@
 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,
     IMPLIES,
+    KEY_REGISTRY,
     LOWER_TRIANGULAR,
     MATRIX_KEYS,
     ORTHOGONAL,
@@ -35,6 +37,7 @@
     register_assumption,
     register_constant_inference,
     register_implies,
+    register_universal_assumption,
 )
 from pytensor.assumptions.specify import (
     SpecifyAssumptions,

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)

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)

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)