diff --git a/pytensor/tensor/rewriting/linalg/solvers.py b/pytensor/tensor/rewriting/linalg/solvers.py index f3439bd3c6..8fe573eca2 100644 --- a/pytensor/tensor/rewriting/linalg/solvers.py +++ b/pytensor/tensor/rewriting/linalg/solvers.py @@ -15,6 +15,7 @@ from pytensor.scan.op import Scan from pytensor.scan.rewriting import scan_seqopt1 from pytensor.tensor.basic import atleast_Nd, split +from pytensor.tensor.blas import BatchedDot from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.elemwise import DimShuffle from pytensor.tensor.linalg.constructors import BlockDiagonal @@ -34,6 +35,7 @@ tridiagonal_lu_factor, tridiagonal_lu_solve, ) +from pytensor.tensor.math import Dot, _matmul from pytensor.tensor.rewriting.basic import ( register_canonicalize, register_specialize, @@ -111,19 +113,20 @@ def batched_vector_b_solve_to_matrix_b_solve(fgraph, node): @register_stabilize -@node_rewriter([blockwise_of(OpPattern(Solve, b_ndim=2))]) +@node_rewriter([blockwise_of(Solve)]) def psd_solve_to_chol_solve(fgraph, node): """Rewrite solve(A, b) → triangular solves via Cholesky when A is positive-definite.""" - assume_a = node.op.core_op.assume_a + core_op = node.op.core_op A, b = node.inputs if ( - assume_a == "pos" + core_op.assume_a == "pos" or getattr(A.tag, "psd", None) is True or check_assumption(fgraph, A, POSITIVE_DEFINITE) ): + b_ndim = core_op.b_ndim L = cholesky(A) - Li_b = solve_triangular(L, b, lower=True, b_ndim=2) - x = solve_triangular((L.mT), Li_b, lower=False, b_ndim=2) + Li_b = solve_triangular(L, b, lower=True, b_ndim=b_ndim) + x = solve_triangular((L.mT), Li_b, lower=False, b_ndim=b_ndim) return [x] @@ -701,3 +704,98 @@ def jax_bilinear_lyapunov_to_direct(fgraph, node): "jax", position=0.9, # Run before canonicalization ) + + +def _peel_shuffles_of_quadratic_form(var): + r"""Strip DimShuffles that leave a quadratic form's reading intact. + + Adding or dropping broadcastable axes is safe, and so is the matrix transpose that + writes the form as :math:`x^\top M x`. Any other permutation of the retained axes + changes which axis the surrounding dot contracts. + """ + while (owner := var.owner) is not None and isinstance(owner.op, DimShuffle): + op = owner.op + axes_kept_in_order = tuple(sorted(op.shuffle)) == op.shuffle + if not (axes_kept_in_order or op.is_left_expanded_matrix_transpose): + break + var = owner.inputs[0] + return var + + +@register_specialize +@node_rewriter([Dot, _matmul, BatchedDot]) +def quadratic_form_to_single_solve(fgraph, node): + r"""Halve the substitutions in the quadratic form :math:`b^\top K^{-1} b`. + + Applying :math:`K^{-1}` to :math:`b` costs a forward *and* a back substitution, but a + surrounding dot against the same :math:`b` contracts the result straight back to a + scalar. Writing the triangular factor as :math:`C`, + + .. math:: + + b^\top K^{-1} b = \lVert C^{-1} b \rVert^2 \quad (K = C C^\top), \\ + b^\top K^{-1} b = \lVert C^{-\top} b \rVert^2 \quad (K = C^\top C), + + so one triangular solve suffices. Matches ``cho_solve`` and a hand-written pair of + transposed triangular solves, whose factor need not come from a Cholesky. + """ + [out] = node.outputs + # A quadratic form contracts both core axes away, leaving them singleton. + if out.type.ndim < 2 or any(dim != 1 for dim in out.type.shape[-2:]): + return None + + # Either input may hold the solve, under the shuffles its variant puts on each side. + left, right = (_peel_shuffles_of_quadratic_form(inp) for inp in node.inputs) + for solve_side, other in ((left, right), (right, left)): + owner = solve_side.owner + if owner is None: + continue + core_op = owner.op + if isinstance(core_op, Blockwise): + core_op = core_op.core_op + + if isinstance(core_op, CholeskySolve): + factor, b = owner.inputs[:2] + if ( + _peel_shuffles_of_quadratic_form(b) is not other + # The cho_solve would stay for its other consumers, making the new + # solve a third substitution rather than a replacement for the second. + or len(fgraph.clients[solve_side]) > 1 + # cho_solve reads a complex K as C C^H, which the plain transpose of + # the surrounding dot does not undo. + or factor.type.dtype.startswith("complex") + ): + continue + lower, b_ndim = core_op.lower, core_op.b_ndim + # An upper factor means K = C.T @ C, so the single solve is against C.T. + z = solve_triangular( + factor, b, lower=lower, trans=0 if lower else 1, b_ndim=b_ndim + ) + break + + if isinstance(core_op, SolveTriangular): + # solve_triangular(A.mT, solve_triangular(A, b, lower=l), lower=not l), + # which is also what the user-facing ``trans`` flag lowers to. The inner + # solve is kept, so the rewrite pays off however many clients it has. + factor_T, inner = owner.inputs + match (factor_T.owner_op_and_inputs, inner.owner_op_and_inputs): + case ( + (DimShuffle(is_left_expanded_matrix_transpose=True), factor), + (Blockwise(SolveTriangular() as inner_op), inner_factor, b), + ) if ( + inner_factor is factor + and inner_op.lower != core_op.lower + and inner_op.unit_diagonal == core_op.unit_diagonal + and _peel_shuffles_of_quadratic_form(b) is other + ): + z, b_ndim = inner, inner_op.b_ndim + break + else: + return None + + core_axes = tuple(range(z.type.ndim - b_ndim, z.type.ndim)) + # Restore the singleton core axes the dot contracted away. Not via ``out.shape``, + # which would make the replacement depend on the node it replaces. + quad = pt.shape_padright((z * z).sum(axis=core_axes), 2) + copy_stack_trace(out, quad) + return [quad] diff --git a/tests/tensor/rewriting/linalg/test_solvers.py b/tests/tensor/rewriting/linalg/test_solvers.py index 3dd70ae394..56a4e49709 100644 --- a/tests/tensor/rewriting/linalg/test_solvers.py +++ b/tests/tensor/rewriting/linalg/test_solvers.py @@ -29,11 +29,12 @@ SolveLUFactorTridiagonal, ) from pytensor.tensor.rewriting.linalg.solvers import ( + quadratic_form_to_single_solve, reuse_decomposition_multiple_solves, scan_split_non_sequence_decomposition_and_solve, ) -from pytensor.tensor.type import matrix, tensor -from tests.unittest_tools import assert_equal_computations +from pytensor.tensor.type import matrix, tensor, vector +from tests.unittest_tools import RewriteTester, assert_equal_computations def test_generic_solve_to_solve_triangular(): @@ -77,17 +78,18 @@ def test_generic_solve_to_solve_triangular(): ) -def test_psd_solve_with_chol(): +@pytest.mark.parametrize("b_ndim", [1, 2]) +def test_psd_solve_with_chol(b_ndim): """Test that solve(A, b) with PSD A gets rewritten to cholesky + cho_solve.""" A = matrix("A") - b = matrix("b") + b = vector("b") if b_ndim == 1 else matrix("b") A_psd = assume(A, positive_definite=True) out = pt.linalg.solve(A_psd, b) rewritten = rewrite_graph(out, include=("canonicalize", "stabilize", "specialize")) L = cholesky(A_psd) - expected = cho_solve((L, True), b, b_ndim=2) + expected = cho_solve((L, True), b, b_ndim=b_ndim) assert_equal_computations([rewritten], [expected]) @@ -662,3 +664,171 @@ def test_orthogonal_solve_to_transpose_matmul(): rewritten = rewrite_graph(out, include=rewrites) expected = rewrite_graph(Q_orth.mT @ b, include=rewrites) assert_equal_computations([rewritten], [expected]) + + +# The rewrite is registered in specialize, and reaching it needs the earlier passes: +# canonicalize lowers the dot, stabilize turns a positive-definite solve into a +# Cholesky pair. +REWRITE_PASSES = ("canonicalize", "stabilize", "specialize") + + +class TestQuadraticFormToSingleSolve: + # Inputs for the table of graphs the rewrite must leave alone. + psd = matrix("psd", shape=(6, 6)) + dense = matrix("dense", shape=(6, 6)) + tri_factor = matrix("tri_factor", shape=(6, 6)) + other_tri_factor = matrix("other_tri_factor", shape=(6, 6)) + lhs, rhs = vector("lhs", shape=(6,)), vector("rhs", shape=(6,)) + complex_psd = matrix("complex_psd", shape=(4, 4), dtype="complex128") + complex_rhs = vector("complex_rhs", shape=(4,), dtype="complex128") + chol = cholesky(psd, lower=True) + cho_solved = cho_solve((chol, True), rhs) + inner_solve = solve_triangular(tri_factor, rhs, lower=True) + + @pytest.mark.parametrize( + "solve_on_left", [False, True], ids=["solve_on_right", "solve_on_left"] + ) + @pytest.mark.parametrize("batch", [(), (4,)], ids=["unbatched", "batched"]) + @pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"]) + def test_cho_solve_variant(self, lower, batch, solve_on_left): + """An upper factor means ``cho_solve`` reads ``K = C.T @ C``, so the single + remaining solve is against ``C.T``.""" + K = tensor("K", shape=(*batch, 6, 6)) + y = tensor("y", shape=(*batch, 6)) + factor = cholesky(K, lower=lower) + solved = cho_solve((factor, lower), y, b_ndim=1) + quad = ( + solved[..., None, :] @ y[..., :, None] + if solve_on_left + else y[..., None, :] @ solved[..., :, None] + ) + + result = RewriteTester([K, y], [quad.squeeze((-2, -1))], include=REWRITE_PASSES) + + z = solve_triangular(factor if lower else factor.mT, y, lower=True, b_ndim=1) + result.assert_graph(pt.sqr(z).sum(axis=-1)) + + rng = np.random.default_rng(0) + M = rng.normal(size=(*batch, 6, 6)) + result.assert_eval( + M @ np.swapaxes(M, -1, -2) + 6 * np.eye(6), rng.normal(size=(*batch, 6)) + ) + + @pytest.mark.parametrize("dtype", ["float64", "complex128"]) + @pytest.mark.parametrize( + "unit_diagonal", [False, True], ids=["stored_diagonal", "unit_diagonal"] + ) + @pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"]) + def test_paired_solves_variant(self, lower, unit_diagonal, dtype): + """``b.T @ A^-T @ A^-1 @ b`` is the squared norm of the inner solve for any + invertible ``A``, Cholesky factor or not, and for any dtype: both transposes are + unconjugated. The inner solve is kept, so another consumer of it does not hold + the rewrite back.""" + C = matrix("C", shape=(6, 6), dtype=dtype) + y = vector("y", shape=(6,), dtype=dtype) + inner = solve_triangular(C, y, lower=lower, unit_diagonal=unit_diagonal) + outer = solve_triangular( + C.mT, inner, lower=not lower, unit_diagonal=unit_diagonal + ) + + result = RewriteTester( + [C, y], [y @ outer, pt.exp(inner)], include=REWRITE_PASSES + ) + + result.assert_graph(pt.sqr(inner).sum(axis=-1), pt.exp(inner)) + + rng = np.random.default_rng(7) + M = rng.normal(size=(6, 6)) + 6 * np.eye(6) + y_val = rng.normal(size=6) + if dtype == "complex128": + M = M + 1j * rng.normal(size=(6, 6)) + y_val = y_val + 1j * rng.normal(size=6) + result.assert_eval(np.tril(M) if lower else np.triu(M), y_val) + + @pytest.mark.parametrize( + "transposed", [False, True], ids=["promoted", "transposed"] + ) + def test_column_right_hand_side(self, transposed): + """A matrix rhs of one column is a quadratic form under DimShuffles: promoting a + vector and indexing back down, or a column already written as ``x.mT @ ...``.""" + K = matrix("K", shape=(6, 6)) + L = cholesky(K, lower=True) + x = matrix("x", shape=(6, 1)) if transposed else vector("x", shape=(6,)) + solved = cho_solve((L, True), x if transposed else x[:, None]) + out = x.mT @ solved if transposed else x @ solved[:, 0] + + result = RewriteTester([K, x], [out], include=REWRITE_PASSES) + + z = solve_triangular( + L, x if transposed else pt.expand_dims(x, 1), lower=True, b_ndim=2 + ).squeeze(1) + quad = pt.sqr(z).sum(axis=-1) + result.assert_graph(pt.expand_dims(quad, (0, 1)) if transposed else quad) + + rng = np.random.default_rng(5) + M = rng.normal(size=(6, 6)) + result.assert_eval( + M @ M.T + 6 * np.eye(6), rng.normal(size=(6, 1) if transposed else 6) + ) + + @pytest.mark.parametrize( + "inputs, outputs", + [ + pytest.param( + [psd, dense], + [dense @ cho_solve((chol, True), dense)], + id="core_axes_survive", + ), + pytest.param( + [psd, lhs, rhs], + [lhs @ cho_solve((chol, True), rhs)], + id="contracted_against_other", + ), + # cho_solve reads a complex K as C @ C.conj().T, which the dot's + # unconjugated transpose does not undo into a squared norm. + pytest.param( + [complex_psd, complex_rhs], + [complex_rhs @ cho_solve((cholesky(complex_psd), True), complex_rhs)], + id="complex_cho_solve", + ), + # The cho_solve stays for its other consumer, so collapsing the quadratic + # form would add a third substitution instead of removing one. + pytest.param( + [psd, rhs], + [rhs @ cho_solved, pt.exp(cho_solved)], + id="reused_cho_solve", + ), + pytest.param( + [tri_factor, other_tri_factor, rhs], + [rhs @ solve_triangular(other_tri_factor.mT, inner_solve, lower=False)], + id="outer_factor_does_not_undo_inner", + ), + pytest.param( + [tri_factor, rhs], + [ + rhs + @ solve_triangular( + tri_factor.mT, inner_solve, lower=False, unit_diagonal=True + ) + ], + id="unit_diagonal_mismatch", + ), + ], + ) + def test_not_applied(self, inputs, outputs): + """Each of these looks like a quadratic form somewhere and is not one, so the + graph has to come out as the other passes left it.""" + rewritten = RewriteTester(inputs, outputs, include=REWRITE_PASSES) + untouched = RewriteTester( + inputs, + outputs, + include=REWRITE_PASSES, + exclude=(quadratic_form_to_single_solve.name,), + ) + + assert_equal_computations( + rewritten.rewr_fg.outputs, + untouched.rewr_fg.outputs, + in_xs=rewritten.rewr_fg.inputs, + in_ys=untouched.rewr_fg.inputs, + )