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
@@ -154,8 +155,26 @@ 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.
     """
+    a = pt.as_tensor_variable(a)
     assume_a = assume_a.lower()
 
+    # Restate what ``assume_a`` promises as an assumption, so every other consumer of ``a`` can
+    # act on it too. "general" promises nothing, and "tridiagonal" and "banded" have no key.
+    match assume_a:
+        case "diagonal":
+            a = assume(a, diagonal=True)
+        case "lower triangular":
+            a = assume(a, lower_triangular=True)
+        case "upper triangular":
+            a = assume(a, upper_triangular=True)
+        case "pos" | "positive definite":
+            a = assume(a, positive_definite=True)
+        case "sym" | "symmetric":
+            a = assume(a, symmetric=True)
+        case "her" | "hermitian" if not a.type.dtype.startswith("complex"):
+            # A real Hermitian matrix is symmetric; a complex one is not.
+            a = assume(a, symmetric=True)
+
     if assume_a in ("lower triangular", "upper triangular"):
         lower = "lower" in assume_a
         return solve_triangular(

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,5 +1,6 @@
 from collections.abc import Sequence
 
+from pytensor.assumptions.specify import assume
 from pytensor.tensor.linalg.decomposition.cholesky import Cholesky
 from pytensor.tensor.linalg.solvers.general import Solve
 from pytensor.xtensor.type import as_xtensor
@@ -94,6 +95,19 @@ def solve(
     else:
         raise ValueError("Solve dims must have length 2 or 3")
 
+    assume_a = assume_a.lower()
+
+    # Restate what ``assume_a`` promises as an assumption, so every other consumer of ``a`` can
+    # act on it too. "general" promises nothing, and "tridiagonal" and "banded" have no key.
+    match assume_a:
+        case "pos" | "positive definite":
+            a = assume(a, positive_definite=True)
+        case "sym" | "symmetric":
+            a = assume(a, symmetric=True)
+        case "her" | "hermitian" if not a.type.dtype.startswith("complex"):
+            # A real Hermitian matrix is symmetric; a complex one is not.
+            a = assume(a, symmetric=True)
+
     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)

tests/tensor/linalg/test_solvers/test_general.py

@@ -6,6 +6,7 @@
 
 from pytensor import function
 from pytensor import tensor as pt
+from pytensor.assumptions.specify import assume
 from pytensor.configdefaults import config
 from pytensor.graph.basic import equal_computations
 from pytensor.tensor import TensorVariable
@@ -206,20 +207,164 @@ def test_solve_gradient(
             lambda A, b: solve_op(A_func(A), b), [A_val, b_val], 3, rng, eps=eps
         )
 
+    @staticmethod
+    def _op_names(fn):
+        return [
+            type(getattr(node.op, "core_op", node.op)).__name__
+            for node in fn.maker.fgraph.apply_nodes
+        ]
+
+    @pytest.mark.skipif(
+        config.mode == "FAST_COMPILE", reason="Consumers rely on rewrites"
+    )
+    @pytest.mark.parametrize("assume_a", ["sym", "pos", "her"])
+    def test_assume_a_records_an_assumption_about_a(self, assume_a):
+        """``assume_a`` promises a property of ``a``, so other readers of ``a`` get it.
+
+        Nothing else in the graph tells ``eig`` that ``a`` is symmetric, and ``pos``
+        reaches ``eigh`` through the implication that positive definite matrices are.
+        """
+        a, b = matrix("a"), matrix("b")
+        w, _ = pt.linalg.eig(a)
+        fn = function([a, b], [solve(a, b, assume_a=assume_a), w])
+
+        op_names = self._op_names(fn)
+        assert "Eigh" in op_names
+        assert "Eig" not in op_names
+        assert "SpecifyAssumptions" not in op_names, (
+            "the marker must not outlive the drain"
+        )
+
+        rng = np.random.default_rng(31)
+        X = rng.normal(size=(6, 6)).astype(config.floatX)
+        a_val = X @ X.T + 6 * np.eye(6, dtype=config.floatX)
+        b_val = rng.normal(size=(6, 2)).astype(config.floatX)
+
+        ATOL = 1e-8 if config.floatX.endswith("64") else 1e-4
+        RTOL = 1e-8 if config.floatX.endswith("64") else 1e-4
+        solved, eigenvalues = fn(a_val, b_val)
+        np.testing.assert_allclose(
+            solved, np.linalg.solve(a_val, b_val), atol=ATOL, rtol=RTOL
+        )
+        np.testing.assert_allclose(
+            np.sort(eigenvalues), np.linalg.eigvalsh(a_val), atol=ATOL, rtol=RTOL
+        )
+
+    @pytest.mark.skipif(
+        config.mode == "FAST_COMPILE", reason="Consumers rely on rewrites"
+    )
+    def test_hermitian_is_not_recorded_as_symmetric_for_complex_input(self):
+        """A complex Hermitian matrix satisfies ``a.conj().T == a``, not ``a.T == a``.
+
+        Recording it as symmetric would license every rewrite that transposes ``a``
+        freely, so the promise stops at the solve for complex dtypes.
+        """
+        a = matrix("a", dtype="complex128")
+        b = matrix("b", dtype="complex128")
+        w, _ = pt.linalg.eig(a)
+        fn = function([a, b], [solve(a, b, assume_a="her"), w])
+
+        op_names = self._op_names(fn)
+        assert "Eig" in op_names
+        assert "Eigh" not in op_names
+
+    @pytest.mark.skipif(
+        config.mode == "FAST_COMPILE", reason="Consumers rely on rewrites"
+    )
+    def test_rewrite_built_solve_records_nothing(self):
+        """``inv_to_solve`` builds a solve from an assumption it has already read.
+
+        Recording it again would leave a marker behind, as rewriting runs long after
+        the pass that resolves them.
+        """
+        X, r = matrix("X"), matrix("r")
+        fn = function([X, r], pt.linalg.inv(assume(X, positive_definite=True)) @ r)
+
+        assert "SpecifyAssumptions" not in self._op_names(fn)
+
+    @pytest.mark.skipif(
+        config.mode == "FAST_COMPILE", reason="Consumers rely on rewrites"
+    )
+    def test_assume_a_diagonal_records_an_assumption_about_a(self):
+        """``assume_a='diagonal'`` lowers to a division, erasing the promise.
+
+        Recording it first keeps the property available to every other reader of ``a``,
+        each of which drops from a dense op to an elementwise one.
+        """
+        a, b, c = matrix("a"), matrix("b"), matrix("c")
+        fn = function(
+            [a, b, c], [solve(a, b, assume_a="diagonal"), a @ c, pt.linalg.det(a)]
+        )
+
+        op_names = self._op_names(fn)
+        assert "Dot" not in op_names
+        assert "Det" not in op_names
+
+        rng = np.random.default_rng(42)
+        a_val = np.diag(rng.normal(size=6) + 5.0).astype(config.floatX)
+        b_val = rng.normal(size=(6, 2)).astype(config.floatX)
+        c_val = rng.normal(size=(6, 3)).astype(config.floatX)
+
+        ATOL = 1e-8 if config.floatX.endswith("64") else 1e-4
+        RTOL = 1e-8 if config.floatX.endswith("64") else 1e-4
+        solved, product, det = fn(a_val, b_val, c_val)
+        np.testing.assert_allclose(
+            solved, np.linalg.solve(a_val, b_val), atol=ATOL, rtol=RTOL
+        )
+        np.testing.assert_allclose(product, a_val @ c_val, atol=ATOL, rtol=RTOL)
+        np.testing.assert_allclose(det, np.linalg.det(a_val), atol=ATOL, rtol=RTOL)
+
+    @pytest.mark.skipif(
+        config.mode == "FAST_COMPILE", reason="Consumers rely on rewrites"
+    )
+    def test_assume_a_reaches_an_untagged_solve_of_the_same_matrix(self):
+        """A second solve that made no promise of its own still picks the property up.
+
+        The two use different right-hand sides so that they cannot simply merge.
+        """
+        a, b, c = matrix("a"), matrix("b"), matrix("c")
+        fn = function([a, b, c], [solve(a, b, assume_a="pos"), solve(a, c)])
+
+        op_names = self._op_names(fn)
+        assert "Solve" not in op_names
+        assert op_names.count("Cholesky") == 1
+        assert op_names.count("CholeskySolve") == 2
+
+    @pytest.mark.skipif(
+        config.mode == "FAST_COMPILE", reason="Consumers rely on rewrites"
+    )
+    def test_no_assumption_recorded_without_a_promise(self):
+        """``assume_a='gen'`` asserts nothing, so nothing is recorded about ``a``."""
+        a, b = matrix("a"), matrix("b")
+        w, _ = pt.linalg.eig(a)
+        fn = function([a, b], [solve(a, b), w])
+
+        op_names = self._op_names(fn)
+        assert "Eig" in op_names
+        assert "Eigh" not in op_names
+        assert "Cholesky" not in op_names
+
     def test_solve_tringular_indirection(self):
+        """The triangular assume_a dispatches to solve_triangular and records itself."""
         a = pt.matrix("a")
         b = pt.vector("b")
 
         indirect = solve(a, b, assume_a="lower triangular")
-        direct = solve_triangular(a, b, lower=True, trans=False)
+        direct = solve_triangular(
+            assume(a, lower_triangular=True), b, lower=True, trans=False
+        )
         assert equal_computations([indirect], [direct])
 
         indirect = solve(a, b, assume_a="upper triangular")
-        direct = solve_triangular(a, b, lower=False, trans=False)
+        direct = solve_triangular(
+            assume(a, upper_triangular=True), b, lower=False, trans=False
+        )
         assert equal_computations([indirect], [direct])
 
         indirect = solve(a, b, assume_a="upper triangular", transposed=True)
-        direct = solve_triangular(a, b, lower=False, trans=True)
+        direct = solve_triangular(
+            assume(a, upper_triangular=True), b, lower=False, trans=True
+        )
         assert equal_computations([indirect], [direct])
 
 

tests/tensor/rewriting/linalg/test_decomposition.py

@@ -642,7 +642,7 @@ def test_eig_to_eigh():
     rewrites = ("canonicalize", "ShapeOpt")
     w_r, v_r = rewrite_graph([w, v], include=rewrites)
 
-    w_expected, v_expected = eigh(x_sym)
+    w_expected, v_expected = eigh(x)
     w_expected = w_expected.astype("complex128")
     v_expected = v_expected.astype("complex128")
     assert_equal_computations([w_r, v_r], [w_expected, v_expected])

tests/tensor/rewriting/linalg/test_solvers.py

@@ -86,7 +86,7 @@ def test_psd_solve_with_chol():
 
     rewritten = rewrite_graph(out, include=("canonicalize", "stabilize", "specialize"))
 
-    L = cholesky(A_psd)
+    L = cholesky(A)
     expected = cho_solve((L, True), b, b_ndim=2)
 
     assert_equal_computations([rewritten], [expected])