Make the assumptions system extensible by downstream libraries

[jessegrabowski]

Downstream libraries can define their own assumption keys now, and both the built-in and downstream ones can be set and queried off the key object.

Declaring and checking assumptions

from pytensor.assumptions import SYMMETRIC, assume

x = assume(x, symmetric=True)   # unchanged
x = SYMMETRIC.assume(x)         # same thing, off the key
SYMMETRIC.holds(x)              # True

holds builds a throwaway FunctionGraph per call; pass one in (SYMMETRIC.holds(x, fgraph)) when asking about several variables of the same graph.

Adding your own assumption

Constructing an AssumptionKey registers it, which is all it takes for assume(), debugprint, and the drain rewrite to pick it up. register_matrix_property_rules(key) hands it the built-in trailing-two-axes propagation; a key with its own algebra registers rules the way the built-in properties do.

# before
from pytensor.tensor.rewriting.assumptions import _KEY_BY_NAME

TIME_VARYING = AssumptionKey("time_varying", "tv")
register_assumption(TIME_VARYING, SpecifyAssumptions)(specify_assumption_rule)
_KEY_BY_NAME[TIME_VARYING.name] = TIME_VARYING  # else the drain rewrite KeyErrors

def declare_time_varying(x):
    return SpecifyAssumptions({TIME_VARYING.name: FactState.TRUE})(x)

def is_time_varying(x):
    fgraph = FunctionGraph(outputs=[x], clone=False)
    return check_assumption(fgraph, x, TIME_VARYING)

# after
TIME_VARYING = AssumptionKey("time_varying", "tv")

x = TIME_VARYING.assume(x)      # or assume(x, time_varying=True)
TIME_VARYING.holds(x)

SpecifyAssumptions takes keys, not names

SpecifyAssumptions stores (AssumptionKey, FactState) pairs rather than (str, FactState), so SpecifyAssumptions({"diagonal": ...}) now raises with a pointer to assume(). Holding a key implies it is registered, which makes an unregistered assumption unrepresentable instead of something the drain rewrite has to check for at compile time. AssumptionKey.__reduce__ keeps that true across pickling, where the default dataclass path would skip __init__ and restore a key with no rules installed.

unique_indices no longer inherits the matrix rules

The trailing-two-axes rules were registered for every key rather than just the matrix ones. An Alloc broadcast therefore claimed unique_indices still held on an array where every entry repeats, and that fact exists to license lifting an index earlier.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB