.progress

Author: ricardoV94

pytensor/tensor/_linalg/solve/rewrites.py

@@ -1,5 +1,10 @@
-from pytensor.graph import node_rewriter
-from pytensor.tensor._linalg.solve.tridiagonal import split_solve_tridiagonal, decompose_of_solve_tridiagonal
+from pytensor.graph.rewriting.basic import node_rewriter
+from pytensor.tensor._linalg.solve.tridiagonal import (
+    decompose_tridiagonals,
+    solve_decomposed_tridiagonal,
+    split_solve_tridiagonal,
+)
+from pytensor.tensor.basic import diagonal
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.rewriting.basic import register_specialize
@@ -10,7 +15,9 @@
 @register_specialize
 @node_rewriter(tracks=[Blockwise])
 def batched_solve_decomposition(fgraph, node):
-    if not(isinstance(node.op.core_op, Solve) and node.op.core_op.assume_a == "tridiagonal"):
+    if not (
+        isinstance(node.op.core_op, Solve) and node.op.core_op.assume_a == "tridiagonal"
+    ):
         return
 
     a, b = node.inputs
@@ -20,8 +27,11 @@ def batched_solve_decomposition(fgraph, node):
     # Check if a is broadcasted in computing the output
     if not any(
         a_bcast and not b_bcast
-        for a_bcast, b_bcast
-        in zip(a.type.broadcastable[:batch_ndim], b.type.broadcastable[:batch_ndim], strict=True)
+        for a_bcast, b_bcast in zip(
+            a.type.broadcastable[:batch_ndim],
+            b.type.broadcastable[:batch_ndim],
+            strict=True,
+        )
     ):
         return
 
@@ -32,7 +42,6 @@ def batched_solve_decomposition(fgraph, node):
 @register_specialize
 @node_rewriter([Blockwise])
 def reuse_lu_decomp_multiple_solves(fgraph, node):
-
     if not isinstance(node.op.core_op, Solve):
         return None
 
@@ -43,32 +52,61 @@ def reuse_lu_decomp_multiple_solves(fgraph, node):
         return None
 
     def find_solve_clients(var):
-        return [
-            cl
-            for cl, idx in fgraph.clients[var]
-            if idx == 0
-            and isinstance(cl.op, Blockwise)
-            and isinstance(cl.op.core_op, Solve)
-            and cl.op.core_op.assume_a == assume_a
-        ]
-
+        clients = []
+        for cl, idx in fgraph.clients[var]:
+            if (
+                idx == 0
+                and isinstance(cl.op, Blockwise)
+                and isinstance(cl.op.core_op, Solve)
+            ):
+                clients.append(cl)
+            elif isinstance(cl.op, DimShuffle) and cl.op.is_left_expand_dims:
+                # If it's a left expand_dims, recurse on the output
+                clients.extend(find_solve_clients(cl.outputs[0]))
+        return clients
 
     [A, _] = node.inputs
     if A.owner is not None and isinstance(A.owner.op, DimShuffle):
-        # FIXME: Don't consider if dimshuffle mixes batch and core dims
+        # If this DimShuffle is more than left expand_dims / matrix transpose
+        # We won't find "clients" again and will exit the rewrite
         [A] = A.owner.inputs
 
-    # Find Solve using A
-    A_solve_clients = [(client, False) for client in find_solve_clients(A)]
+    # Find Solve using A (or left expand_dims of A)
+    A_direct_solve_clients = find_solve_clients(A)
 
     # Find Solves using A.T
+    A_transpose_solve_clients = []
     for cl, _ in fgraph.clients[A]:
         if isinstance(cl.op, DimShuffle) and is_matrix_transpose(cl.out):
             A_T = cl.out
-            A_solve_clients.extend((client, True) for client in find_solve_clients(A_T))
+            A_transpose_solve_clients.extend(find_solve_clients(A_T))
+
+    if (len(A_direct_solve_clients) + len(A_transpose_solve_clients)) <= 1:
+        # If we only have one client, we don't need to do anything
+        # It could still be useful to make the diagonal extraction symbolically
+        # In case they are being set in the graph as well (therefore unnecessary)
+        return None
+
+    dl, d, du = (diagonal(A, offset=o, axis1=-2, axis2=-1) for o in (-1, 0, 1))
 
-    A_decomp = decompose_of_solve_tridiagonal(A)
     replacements = {}
-    for client, transpose in A_solve_clients:
-        _, b = client.inputs
-    return replacements
+    if A_direct_solve_clients:
+        A_direct_decomp = decompose_tridiagonals(dl, d, du)
+        for client in A_direct_solve_clients:
+            _, b = client.inputs
+            b_ndim = client.op.core_op.b_ndim
+            replacements[client.outputs[0]] = solve_decomposed_tridiagonal(
+                A_direct_decomp, b, b_ndim=b_ndim
+            )
+
+    if A_transpose_solve_clients:
+        # We just need to swap the off-diagonals
+        A_transpose_decomp = decompose_tridiagonals(du, d, dl)
+        for client in A_transpose_solve_clients:
+            _, b = client.inputs
+            b_ndim = client.op.core_op.b_ndim
+            replacements[client.outputs[0]] = solve_decomposed_tridiagonal(
+                A_transpose_decomp, b, b_ndim=b_ndim
+            )
+
+    return replacements

pytensor/tensor/_linalg/solve/tridiagonal.py

@@ -1,17 +1,22 @@
-import scipy
 import numpy as np
+import scipy
 from scipy.linalg import get_lapack_funcs
 
-from pytensor.graph import Op, Apply
+from pytensor.graph import Apply, Op
 from pytensor.tensor.basic import as_tensor, diagonal
-from pytensor.tensor.type import tensor, vector
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.slinalg import Solve
+from pytensor.tensor.type import tensor, vector
 
 
 class LUFactorTridiagonal(Op):
     """Compute LU factorization of a tridiagonal matrix (lapack gttrf)"""
-    __props__ = ("overwrite_dl", "overwrite_d", "overwrite_du",)
+
+    __props__ = (
+        "overwrite_dl",
+        "overwrite_d",
+        "overwrite_du",
+    )
     gufunc_signature = "(dl),(d),(dl)->(dl),(d),(dl),(du2),(d)"
 
     def __init__(self, overwrite_dl=False, overwrite_d=False, overwrite_du=False):
@@ -29,11 +34,8 @@ def make_node(self, dl, d, du):
         ndl, nd, ndu = (inp.type.shape[-1] for inp in (dl, d, du))
         n = (
             ndl + 1
-            if ndl is not None else (
-                nd if nd is not None else (
-                    ndu + 1 if ndu is not None else None
-                )
-            )
+            if ndl is not None
+            else (nd if nd is not None else (ndu + 1 if ndu is not None else None))
         )
         dummy_arrays = [np.zeros((), dtype=inp.type.dtype) for inp in (dl, d, du)]
         out_dtype = get_lapack_funcs("gttrf", dummy_arrays).dtype
@@ -63,6 +65,7 @@ def perform(self, node, inputs, output_storage):
 
 class SolveLUFactorTridiagonal(Op):
     """Solve a system of linear equations with a tridiagonal coefficient matrix."""
+
     __props__ = ("b_ndim", "overwrite_b")
 
     def __init__(self, b_ndim: int, overwrite_b=False):
@@ -84,21 +87,30 @@ def make_node(self, dl, d, du, du2, ipiv, b):
         if not all(inp.type.ndim == 1 for inp in (dl, d, du, du2, ipiv)):
             raise ValueError("Inputs must be vectors")
 
-        ndl, nd, ndu, ndu2, nipiv = (inp.type.shape[-1] for inp in (dl, d, du, du2, ipiv))
+        ndl, nd, ndu, ndu2, nipiv = (
+            inp.type.shape[-1] for inp in (dl, d, du, du2, ipiv)
+        )
         nb = b.type.shape[0]
         n = (
             ndl + 1
-            if ndl is not None else (
-                nd if nd is not None else (
-                    ndu + 1 if ndu is not None else (
-                        ndu2 + 2 if ndu2 is not None else (
-                            nipiv if nipiv is not None else nb
-                        )
+            if ndl is not None
+            else (
+                nd
+                if nd is not None
+                else (
+                    ndu + 1
+                    if ndu is not None
+                    else (
+                        ndu2 + 2
+                        if ndu2 is not None
+                        else (nipiv if nipiv is not None else nb)
                     )
                 )
             )
         )
-        dummy_arrays = [np.zeros((), dtype=inp.type.dtype) for inp in (dl, d, du, du2, ipiv)]
+        dummy_arrays = [
+            np.zeros((), dtype=inp.type.dtype) for inp in (dl, d, du, du2, ipiv)
+        ]
         # Seems to always be float64?
         out_dtype = get_lapack_funcs("gttrs", dummy_arrays).dtype
         if self.b_ndim == 1:
@@ -111,14 +123,13 @@ def make_node(self, dl, d, du, du2, ipiv, b):
 
     def perform(self, node, inputs, output_storage):
         gttrs = get_lapack_funcs("gttrs", dtype=node.outputs[0].type.dtype)
-        x, _ = gttrs(
-            *inputs, overwrite_b=self.overwrite_b
-        )
+        x, _ = gttrs(*inputs, overwrite_b=self.overwrite_b)
         output_storage[0][0] = x
 
 
 class SolveTridiagonal(Op):
     """Solve a system of linear equations with a tridiagonal dense matrix."""
+
     __props__ = ("b_ndim", "overwrite_b")
 
     def __init__(self, *, b_ndim: int, overwrite_b: bool = False):
@@ -141,7 +152,9 @@ def make_node(self, dl, d, du, b):
             raise TypeError("Diagonals must have the same dtype")
 
         if b.type.ndim != self.b_ndim:
-            raise ValueError(f"Number of dimensions of b does not match promised {self.b_ndim}")
+            raise ValueError(
+                f"Number of dimensions of b does not match promised {self.b_ndim}"
+            )
 
         out_dtype = scipy.linalg.solve(
             np.eye((3), dtype=d.type.dtype),
@@ -156,13 +169,14 @@ def L_op(self, node, inputs, outputs, output_grads):
 
     def perform(self, node, inputs, output_storage):
         [dl, d, du, b] = inputs
-        _gttrf, _gttrs = get_lapack_funcs(('gttrf', 'gttrs'), dtype=node.outputs[0].type.dtype)
+        _gttrf, _gttrs = get_lapack_funcs(
+            ("gttrf", "gttrs"), dtype=node.outputs[0].type.dtype
+        )
 
         dl, d, du, du2, ipiv, _ = _gttrf(dl, d, du)
         x, _ = _gttrs(dl, d, du, du2, ipiv, b, overwrite_b=self.overwrite_b)
         output_storage[0][0] = x
 
-
     def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
         if 3 not in allowed_inplace_inputs:
             return self
@@ -186,6 +200,7 @@ def solve_tridiagonal_from_full_A_b(a, b, b_ndim: int, transposed: bool):
     dl, d, du = (diagonal(a, offset=o, axis1=-2, axis2=-1) for o in (-1, 0, 1))
     return Blockwise(SolveTridiagonal(b_ndim=b_ndim))(dl, d, du)
 
+
 def split_solve_tridiagonal(node):
     """Split a generic solve tridiagonal system into the 3 atomic steps:
        1. Diagonal extractions
@@ -198,11 +213,21 @@ def split_solve_tridiagonal(node):
     core_op = node.op.core_op
     assert isinstance(core_op, Solve) and core_op.assume_a == "tridiagonal"
     a, b = node.inputs
-    dl, d, du, du2, ipiv = decompose_of_solve_tridiagonal(a)
-    return Blockwise(SolveLUFactorTridiagonal(b_ndim=node.op.core_op.b_ndim))(dl, d, du, du2, ipiv, b)
+    a_decomp = decompose_of_solve_tridiagonal(a)
+    return solve_decomposed_tridiagonal(a_decomp, b, b_ndim=core_op.b_ndim)
+
 
 def decompose_of_solve_tridiagonal(a):
     # Return the decomposition of A implied by a solve tridiagonal
     dl, d, du = (diagonal(a, offset=o, axis1=-2, axis2=-1) for o in (-1, 0, 1))
     dl, d, du, du2, ipiv = Blockwise(LUFactorTridiagonal())(dl, d, du)
     return dl, d, du, du2, ipiv
+
+
+def decompose_tridiagonals(dl, d, du):
+    return Blockwise(LUFactorTridiagonal())(dl, d, du)
+
+
+def solve_decomposed_tridiagonal(a_diagonals, b, *, b_ndim: int):
+    dl, d, du, du2, ipiv = a_diagonals
+    return Blockwise(SolveLUFactorTridiagonal(b_ndim=b_ndim))(dl, d, du, du2, ipiv, b)

pytensor/tensor/rewriting/linalg.py

@@ -75,6 +75,13 @@ def is_matrix_transpose(x: TensorVariable) -> bool:
         if ndims < 2:
             return False
         transpose_order = (*range(ndims - 2), ndims - 1, ndims - 2)
+
+        # Allow expand_dims on the left of the transpose
+        if (diff := len(transpose_order) - len(node.op.new_order)) > 0:
+            transpose_order = (
+                *(["x"] * diff),
+                *transpose_order,
+            )
         return node.op.new_order == transpose_order
     return False