From d7071499dd12e3fc23df216a5989e9ca03848ecc Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 22:45:33 -0400 Subject: [PATCH 1/4] Drain SpecifyAssumptions markers wherever they appear The whole-graph pass at 0.8 only sees markers a user built, so it stays as the fast path; the local rewrite covers any a later rewrite introduces. --- pytensor/tensor/rewriting/assumptions.py | 65 +++++++++++++------ tests/assumptions/test_specify.py | 33 ++++++++++ .../rewriting/linalg/test_decomposition.py | 2 +- tests/tensor/rewriting/linalg/test_solvers.py | 2 +- 4 files changed, 79 insertions(+), 23 deletions(-) diff --git a/pytensor/tensor/rewriting/assumptions.py b/pytensor/tensor/rewriting/assumptions.py index a60b7a798d..8e71a356ac 100644 --- a/pytensor/tensor/rewriting/assumptions.py +++ b/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 = 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(), diff --git a/tests/assumptions/test_specify.py b/tests/assumptions/test_specify.py index 0847f48bd5..b8c5d07d7b 100644 --- a/tests/assumptions/test_specify.py +++ b/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) diff --git a/tests/tensor/rewriting/linalg/test_decomposition.py b/tests/tensor/rewriting/linalg/test_decomposition.py index ce8f90f184..6b2f6422f0 100644 --- a/tests/tensor/rewriting/linalg/test_decomposition.py +++ b/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]) diff --git a/tests/tensor/rewriting/linalg/test_solvers.py b/tests/tensor/rewriting/linalg/test_solvers.py index 3dd70ae394..6e7f90e2a1 100644 --- a/tests/tensor/rewriting/linalg/test_solvers.py +++ b/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]) From 12be4cbea1c80385c8f01ca4dbba28adfd7dc386 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 22:47:26 -0400 Subject: [PATCH 2/4] Extract assume_a normalization from Solve.__init__ --- pytensor/tensor/linalg/solvers/general.py | 29 ++++++++++++++--------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/pytensor/tensor/linalg/solvers/general.py b/pytensor/tensor/linalg/solvers/general.py index f7629a6eda..d29cb2ab11 100644 --- a/pytensor/tensor/linalg/solvers/general.py +++ b/pytensor/tensor/linalg/solvers/general.py @@ -14,6 +14,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 +46,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"): @@ -154,7 +161,7 @@ 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() + assume_a = _normalize_assume_a(assume_a) if assume_a in ("lower triangular", "upper triangular"): lower = "lower" in assume_a From bbd38022830264ac4fd9d842b553898f62ee8ead Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 22:48:36 -0400 Subject: [PATCH 3/4] Record assume_a as an assumption about the solved matrix --- pytensor/tensor/linalg/solvers/general.py | 29 ++++ pytensor/xtensor/linalg.py | 8 +- .../linalg/test_solvers/test_general.py | 151 +++++++++++++++++- 3 files changed, 184 insertions(+), 4 deletions(-) diff --git a/pytensor/tensor/linalg/solvers/general.py b/pytensor/tensor/linalg/solvers/general.py index d29cb2ab11..702ca11a91 100644 --- a/pytensor/tensor/linalg/solvers/general.py +++ b/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 @@ -97,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, @@ -161,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. """ + # _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 diff --git a/pytensor/xtensor/linalg.py b/pytensor/xtensor/linalg.py index 05e915f75d..c12eea47fe 100644 --- a/pytensor/xtensor/linalg.py +++ b/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, diff --git a/tests/tensor/linalg/test_solvers/test_general.py b/tests/tensor/linalg/test_solvers/test_general.py index 6c6ff1917b..64baf693ff 100644 --- a/tests/tensor/linalg/test_solvers/test_general.py +++ b/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]) From c7a12cf931d1c10f7f1e7996becbb8dc4108fc24 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 23 Aug 2026 23:30:38 -0400 Subject: [PATCH 4/4] Mypy :D --- pytensor/tensor/rewriting/assumptions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/assumptions.py b/pytensor/tensor/rewriting/assumptions.py index 8e71a356ac..d5e2662c10 100644 --- a/pytensor/tensor/rewriting/assumptions.py +++ b/pytensor/tensor/rewriting/assumptions.py @@ -30,7 +30,7 @@ def _drain_marker(feature: AssumptionFeature, node) -> Variable: for name, _ in node.op.assumptions: feature.get(out, _KEY_BY_NAME[name]) - inp = node.inputs[0] + inp: Variable = node.inputs[0] while inp.owner is not None and isinstance(inp.owner.op, SpecifyAssumptions): inp = inp.owner.inputs[0] return inp