Collapse quadratic forms to a single triangular solve

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]

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,
+        )