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,118 @@ def jax_bilinear_lyapunov_to_direct(fgraph, node):
     "jax",
     position=0.9,  # Run before canonicalization
 )
+
+
+def _peel_shuffles_of_quadratic_form(var):
+    """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 and stops the peel.
+    """
+    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
+
+
+def _single_solve_of_quadratic_form(var):
+    r"""Recognize ``var`` as :math:`A^{-\top} A^{-1} b`, applied by either spelling.
+
+    Both ``cho_solve`` and an explicit pair of transposed triangular solves apply that
+    operator, so a dot against ``b`` collapses to :math:`\lVert z \rVert^2`. Return
+    ``None`` when ``var`` applies neither spelling.
+
+    Returns
+    -------
+    z : TensorVariable
+        The single solve :math:`A^{-1} b`.
+    b : TensorVariable
+        The right-hand side the dot must contract against for the collapse to hold.
+    b_ndim : int
+        Core rank of ``b``, 1 for a vector and 2 for a matrix.
+    """
+    owner = var.owner
+    if owner is None:
+        return None
+
+    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]
+        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
+        )
+        return z, b, b_ndim
+
+    if isinstance(core_op, SolveTriangular):
+        # solve_triangular(A.mT, solve_triangular(A, b, lower=l), lower=not l). The
+        # user-facing ``trans`` flag lowers to exactly this, so both spellings match.
+        factor_T, inner = owner.inputs
+        match factor_T.owner_op_and_inputs:
+            case (DimShuffle(is_left_expanded_matrix_transpose=True), factor):
+                pass
+            case _:
+                return None
+        match inner.owner_op_and_inputs:
+            case (Blockwise(SolveTriangular() as inner_core_op), inner_factor, b) if (
+                inner_factor is factor
+                and inner_core_op.lower != core_op.lower
+                and inner_core_op.unit_diagonal == core_op.unit_diagonal
+            ):
+                return inner, b, inner_core_op.b_ndim
+
+    return None
+
+
+@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 a single triangular solve suffices. This catches both ``cho_solve`` -- transposed
+    when the factor is upper, since ``cho_solve`` then factors :math:`K` as
+    :math:`C^\top C` -- 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
+
+    left, right = (_peel_shuffles_of_quadratic_form(inp) for inp in node.inputs)
+    for solve_side, other in ((left, right), (right, left)):
+        found = _single_solve_of_quadratic_form(solve_side)
+        if found is None:
+            continue
+        z, b, b_ndim = found
+        # Soundness: the vector contracted against must be the solve's own rhs.
+        # Compare through the broadcast shuffles each side may carry.
+        if _peel_shuffles_of_quadratic_form(b) is other:
+            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. This must not be
+    # expressed via ``out.shape``: referencing the replaced node's own shape 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]