Record assume_a as an assumption about the solved matrix

pytensor/tensor/linalg/solvers/general.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from pytensor import tensor as pt
+from pytensor.assumptions.specify import assume
 from pytensor.graph.op import Op
 from pytensor.tensor.basic import diagonal
 from pytensor.tensor.blockwise import Blockwise
@@ -14,6 +15,21 @@
 from pytensor.tensor.variable import TensorVariable
 
 
+# ``Solve`` uses the short spellings, as the various backend dispatches are more
+# likely to recognize them.
+_ASSUME_A_LONG_TO_SHORT = {
+    "general": "gen",
+    "symmetric": "sym",
+    "hermitian": "her",
+    "positive definite": "pos",
+}
+
+
+def _normalize_assume_a(assume_a: str) -> str:
+    assume_a = assume_a.lower()
+    return _ASSUME_A_LONG_TO_SHORT.get(assume_a, assume_a)
+
+
 class Solve(SolveBase):
     """
     Solve a system of linear equations.
@@ -31,19 +47,11 @@ def __init__(self, *, assume_a="gen", **kwargs):
         # Triangular and diagonal are handled outside of Solve
         valid_options = ["gen", "sym", "her", "pos", "tridiagonal", "banded"]
 
-        assume_a = assume_a.lower()
-        # We use the old names as the different dispatches are more likely to support them
-        long_to_short = {
-            "general": "gen",
-            "symmetric": "sym",
-            "hermitian": "her",
-            "positive definite": "pos",
-        }
-        assume_a = long_to_short.get(assume_a, assume_a)
+        assume_a = _normalize_assume_a(assume_a)
 
         if assume_a not in valid_options:
             raise ValueError(
-                f"Invalid assume_a: {assume_a}. It must be one of {valid_options} or {list(long_to_short.keys())}"
+                f"Invalid assume_a: {assume_a}. It must be one of {valid_options} or {list(_ASSUME_A_LONG_TO_SHORT)}"
             )
 
         if assume_a in ("tridiagonal", "banded"):
@@ -90,6 +98,31 @@ def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
         return type(self)(**new_props)
 
 
+def _record_assume_a(a, assume_a: str):
+    """Restate the structure ``assume_a`` promises as an assumption about ``a``.
+
+    ``assume_a`` and :func:`assume` are the same promise in two spellings, but only the
+    solve can act on the first. Recording it lets every other consumer of ``a`` use it.
+    """
+    match assume_a:
+        case "diagonal":
+            return assume(a, diagonal=True)
+        case "lower triangular":
+            return assume(a, lower_triangular=True)
+        case "upper triangular":
+            return assume(a, upper_triangular=True)
+        case "pos":
+            return assume(a, positive_definite=True)
+        case "sym":
+            return assume(a, symmetric=True)
+        case "her" if not a.type.dtype.startswith("complex"):
+            # A real Hermitian matrix is symmetric; a complex one is not.
+            return assume(a, symmetric=True)
+        case _:
+            # "gen" promises nothing, and "tridiagonal" and "banded" have no key.
+            return a
+
+
 def solve(
     a,
     b,
@@ -154,7 +187,10 @@ def solve(
         This will influence how batched dimensions are interpreted.
         By default, we assume b_ndim = b.ndim is 2 if b.ndim > 1, else 1.
     """
-    assume_a = assume_a.lower()
+    # _record_assume_a reads the dtype off ``a``, so it has to be a variable first.
+    a = pt.as_tensor_variable(a)
+    assume_a = _normalize_assume_a(assume_a)
+    a = _record_assume_a(a, assume_a)
 
     if assume_a in ("lower triangular", "upper triangular"):
         lower = "lower" in assume_a

pytensor/tensor/rewriting/assumptions.py

@@ -1,12 +1,41 @@
 from pytensor.assumptions import ALL_KEYS, AssumptionFeature
 from pytensor.assumptions.specify import SpecifyAssumptions
 from pytensor.compile.mode import optdb
-from pytensor.graph.rewriting.basic import GraphRewriter
+from pytensor.graph.basic import Variable
+from pytensor.graph.rewriting.basic import GraphRewriter, node_rewriter
+from pytensor.tensor.rewriting.basic import (
+    register_canonicalize,
+    register_specialize,
+    register_stabilize,
+)
 
 
 _KEY_BY_NAME = {key.name: key for key in ALL_KEYS}
 
 
+def _assumption_feature(fgraph) -> AssumptionFeature:
+    feature = getattr(fgraph, "assumption_feature", None)
+    if feature is None:
+        feature = AssumptionFeature()
+        fgraph.attach_feature(feature)
+    return feature
+
+
+def _drain_marker(feature: AssumptionFeature, node) -> Variable:
+    """Resolve one marker's declarations, returning the input to redirect consumers to.
+
+    Nested markers are peeled so that ``assume(assume(...))`` collapses in one step.
+    """
+    [out] = node.outputs
+    for name, _ in node.op.assumptions:
+        feature.get(out, _KEY_BY_NAME[name])
+
+    inp: Variable = node.inputs[0]
+    while inp.owner is not None and isinstance(inp.owner.op, SpecifyAssumptions):
+        inp = inp.owner.inputs[0]
+    return inp
+
+
 class DrainSpecifyAssumptions(GraphRewriter):
     """Drain ``SpecifyAssumptions`` declarations into the ``AssumptionFeature`` and
     remove the marker nodes.
@@ -29,32 +58,26 @@ def apply(self, fgraph):
             if isinstance(node.op, SpecifyAssumptions)
         ]
 
-        assumption_feature = getattr(fgraph, "assumption_feature", None)
-        if assumption_feature is None:
-            assumption_feature = AssumptionFeature()
-            fgraph.attach_feature(assumption_feature)
-
-        replacements = {}
-        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_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.
-            inp = node.inputs[0]
-            while inp.owner is not None and isinstance(
-                inp.owner.op, SpecifyAssumptions
-            ):
-                inp = inp.owner.inputs[0]
-            replacements[out] = inp
+        feature = _assumption_feature(fgraph)
+        replacements = {node.outputs[0]: _drain_marker(feature, node) for node in nodes}
 
         fgraph.replace_all(
             tuple(replacements.items()), reason="drain_specify_assumptions"
         )
 
 
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([SpecifyAssumptions])
+def drain_specify_assumptions_node(fgraph, node):
+    """Drain a marker that appears after the whole-graph pass has already run.
+
+    A rewrite can then declare an assumption the same way construction does.
+    """
+    return [_drain_marker(_assumption_feature(fgraph), node)]
+
+
 optdb.register(
     "drain_specify_assumptions",
     DrainSpecifyAssumptions(),

pytensor/xtensor/linalg.py

@@ -1,7 +1,11 @@
 from collections.abc import Sequence
 
 from pytensor.tensor.linalg.decomposition.cholesky import Cholesky
-from pytensor.tensor.linalg.solvers.general import Solve
+from pytensor.tensor.linalg.solvers.general import (
+    Solve,
+    _normalize_assume_a,
+    _record_assume_a,
+)
 from pytensor.xtensor.type import as_xtensor
 from pytensor.xtensor.vectorization import XBlockwise
 
@@ -94,6 +98,8 @@ def solve(
     else:
         raise ValueError("Solve dims must have length 2 or 3")
 
+    assume_a = _normalize_assume_a(assume_a)
+    a = _record_assume_a(a, assume_a)
     core_op = Solve(b_ndim=b_ndim, assume_a=assume_a, lower=lower)
     x_op = XBlockwise(
         core_op,

tests/assumptions/test_specify.py

@@ -13,6 +13,10 @@
     FactState,
 )
 from pytensor.assumptions.specify import SpecifyAssumptions, assume
+from pytensor.tensor.rewriting.assumptions import (
+    DrainSpecifyAssumptions,
+    drain_specify_assumptions_node,
+)
 from tests.assumptions.conftest import make_fgraph
 
 
@@ -94,3 +98,32 @@ def test_assume_conflict_with_inferred_fact_raises():
     _, af = make_fgraph(e_not_diag)
     with pytest.raises(ConflictingAssumptionsError):
         af.get(e_not_diag, DIAGONAL)
+
+
+def test_whole_graph_drain_moves_the_fact_onto_the_input():
+    """Draining is not just node removal; the declaration has to land on ``x`` itself.
+
+    Consumers query ``x``, never the marker, so a drain that drops the node without
+    transferring the fact discards the assumption without any visible failure.
+    """
+    x = pt.matrix("x")
+    out = pt.linalg.det(assume(x, positive_definite=True))
+    fg, feature = make_fgraph(out)
+
+    DrainSpecifyAssumptions().apply(fg)
+
+    assert not any(isinstance(node.op, SpecifyAssumptions) for node in fg.apply_nodes)
+    assert feature.check(x, POSITIVE_DEFINITE)
+
+
+def test_local_drain_moves_the_fact_onto_the_input():
+    """The per-node drain owes the same guarantee as the whole-graph pass."""
+    x = pt.matrix("x")
+    marker = assume(x, positive_definite=True)
+    fg, feature = make_fgraph(pt.linalg.det(marker))
+
+    [replacement] = drain_specify_assumptions_node.transform(fg, marker.owner)
+    fg.replace(marker, replacement, reason="test")
+
+    assert replacement is x
+    assert feature.check(x, POSITIVE_DEFINITE)