.progress

Author: ricardoV94

pytensor/scan/rewriting.py

@@ -2561,14 +2561,14 @@ def scan_push_out_dot1(fgraph, node):
     position=1,
 )
 
-scan_seqopt1.register(
-    "scan_decompose_compound_ops,
-    in2out(scan_decompose_compound_ops),
-    "fast_run",
-    "scan",
-    "scan_pushout",
-    position=2,
-)
+# scan_seqopt1.register(
+#     "scan_decompose_compound_ops",
+#     in2out(scan_decompose_compound_ops),
+#     "fast_run",
+#     "scan",
+#     "scan_pushout",
+#     position=2,
+# )
 
 scan_seqopt1.register(
     "scan_push_out_non_seq",

pytensor/tensor/__init__.py

@@ -114,6 +114,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 
 
 # isort: off
+import pytensor.tensor._linalg
 from pytensor.tensor import linalg
 from pytensor.tensor import special
 from pytensor.tensor import signal
@@ -143,6 +144,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
     specify_shape,
 )
 
+
 # We import as `_shared` instead of `shared` to avoid confusion between
 # `pytensor.shared` and `tensor._shared`.
 from pytensor.tensor.sort import argsort, sort

pytensor/tensor/_linalg/__init__.py

@@ -0,0 +1,2 @@
+# Register rewrites
+import pytensor.tensor._linalg.solve

pytensor/tensor/_linalg/solve/__init__.py

@@ -0,0 +1,2 @@
+# Register rewrites in the database
+import pytensor.tensor._linalg.solve.rewrites

pytensor/tensor/_linalg/solve/rewrites.py

@@ -0,0 +1,74 @@
+from pytensor.graph import node_rewriter
+from pytensor.tensor._linalg.solve.tridiagonal import split_solve_tridiagonal, decompose_of_solve_tridiagonal
+from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.rewriting.basic import register_specialize
+from pytensor.tensor.rewriting.linalg import is_matrix_transpose
+from pytensor.tensor.slinalg import Solve
+
+
+@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"):
+        return
+
+    a, b = node.inputs
+    [out] = node.outputs
+    batch_ndim = node.op.batch_ndim(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)
+    ):
+        return
+
+    new_out = split_solve_tridiagonal(node)
+    return [new_out]
+
+
+@register_specialize
+@node_rewriter([Blockwise])
+def reuse_lu_decomp_multiple_solves(fgraph, node):
+
+    if not isinstance(node.op.core_op, Solve):
+        return None
+
+    assume_a = node.op.core_op.assume_a
+
+    if assume_a != "tridiagonal":
+        # Other assume_a not yet supported
+        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
+        ]
+
+
+    [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
+        [A] = A.owner.inputs
+
+    # Find Solve using A
+    A_solve_clients = [(client, False) for client in find_solve_clients(A)]
+
+    # Find Solves using A.T
+    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_decomp = decompose_of_solve_tridiagonal(A)
+    replacements = {}
+    for client, transpose in A_solve_clients:
+        _, b = client.inputs
+    return replacements

pytensor/tensor/_linalg/solve/tridiagonal.py

@@ -3,14 +3,16 @@
 from scipy.linalg import get_lapack_funcs
 
 from pytensor.graph import Op, Apply
-from pytensor.tensor import as_tensor, tensor, diagonal
+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
 
 
 class LUFactorTridiagonal(Op):
     """Compute LU factorization of a tridiagonal matrix (lapack gttrf)"""
     __props__ = ("overwrite_dl", "overwrite_d", "overwrite_du",)
-    _gufunc_signature = "(dl),(d),(dl)->(dl),(d),(dl),(du2),(d)"
+    gufunc_signature = "(dl),(d),(dl)->(dl),(d),(dl),(du2),(d)"
 
     def __init__(self, overwrite_dl=False, overwrite_d=False, overwrite_du=False):
         self.overwrite_dl = overwrite_dl
@@ -19,33 +21,34 @@ def __init__(self, overwrite_dl=False, overwrite_d=False, overwrite_du=False):
         super().__init__()
 
     def make_node(self, dl, d, du):
-        dl, d, du = map(as_tensor, dl, d, du)
+        dl, d, du = map(as_tensor, (dl, d, du))
 
-        if not all(inp.type.ndim == 1 for inp in (dl, d, du))
+        if not all(inp.type.ndim == 1 for inp in (dl, d, du)):
             raise ValueError("Diagonals must be vectors")
 
         ndl, nd, ndu = (inp.type.shape[-1] for inp in (dl, d, du))
         n = (
             ndl + 1
             if ndl is not None else (
-                n if n is not None else (
-                    ndu + 1 if nu is not None else None
+                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
         outputs = [
-            vector(shape=(shape=(None if n is None else n - 1,), dtype=out_dtype),
+            vector(shape=(None if n is None else (n - 1),), dtype=out_dtype),
             vector(shape=(n,), dtype=out_dtype),
             vector(shape=(None if n is None else n - 1,), dtype=out_dtype),
             vector(shape=(None if n is None else n - 2,), dtype=out_dtype),
             vector(shape=(n,), dtype=np.int32),
         ]
+        return Apply(self, [dl, d, du], outputs)
 
     def perform(self, node, inputs, output_storage):
         gttrf = get_lapack_funcs("gttrf", dtype=node.outputs[0].type.dtype)
-        dl, d, du, du2, ipiv, _ = _gttrf(
+        dl, d, du, du2, ipiv, _ = gttrf(
             *inputs,
             overwrite_dl=self.overwrite_dl,
             overwrite_d=self.overwrite_d,
@@ -68,26 +71,26 @@ def __init__(self, b_ndim: int, overwrite_b=False):
         self.b_ndim = b_ndim
         self.overwrite_b = overwrite_b
         if b_ndim == 1:
-            _gufunc_signature = "(dl),(d),(dl),(du2),(d),(d)->(d)
+            self.gufunc_signature = "(dl),(d),(dl),(du2),(d),(d)->(d)"
         else:
-            _gufunc_signature = "(dl),(d),(dl),(du2),(d),(d,rhs)->(d,rhs)
+            self.gufunc_signature = "(dl),(d),(dl),(du2),(d),(d,rhs)->(d,rhs)"
 
     def make_node(self, dl, d, du, du2, ipiv, b):
-        dl, d, du, du2, ipiv, b = map(as_tensor, dl, d, du, du2, ipiv, b)
+        dl, d, du, du2, ipiv, b = map(as_tensor, (dl, d, du, du2, ipiv, b))
 
         if b.type.ndim != self.b_ndim:
             raise ValueError("Wrang number of dimensions for input b.")
 
-        if not all(inp.type.ndim == 1 for inp in (dl, d, du, du2, ipiv))
+        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))
         nb = b.type.shape[0]
         n = (
             ndl + 1
             if ndl is not None else (
-                n if n is not None else (
-                    ndu + 1 if nu 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
                         )
@@ -101,14 +104,14 @@ def make_node(self, dl, d, du, du2, ipiv, b):
         if self.b_ndim == 1:
             output_shape = (n,)
         else:
-            output_shape = (n, n.type.shape[-1])
+            output_shape = (n, b.type.shape[-1])
 
-        outputs = [vector(shape=output_shape, dtype=out_dtype)]
+        outputs = [tensor(shape=output_shape, dtype=out_dtype)]
         return Apply(self, [dl, d, du, du2, ipiv, b], outputs)
 
     def perform(self, node, inputs, output_storage):
         gttrs = get_lapack_funcs("gttrs", dtype=node.outputs[0].type.dtype)
-        x, _ = _gttrs(
+        x, _ = gttrs(
             *inputs, overwrite_b=self.overwrite_b
         )
         output_storage[0][0] = x
@@ -149,7 +152,7 @@ def make_node(self, dl, d, du, b):
         return Apply(self, [dl, d, du, b], [out])
 
     def L_op(self, node, inputs, outputs, output_grads):
-        # TODO
+        pass
 
     def perform(self, node, inputs, output_storage):
         [dl, d, du, b] = inputs
@@ -193,8 +196,13 @@ def split_solve_tridiagonal(node):
     """
     assert isinstance(node.op, Blockwise)
     core_op = node.op.core_op
-    assert isinstance(core_op, Solve) and core.op.assume_a == "tridiagonal"
+    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)
+
+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 Blockwise(SolveLUFactorTridiagonal(b_ndim=node.op.core.op.b_ndim))(dl, d, du)(dl, d, du, du2, ipiv)
+    return dl, d, du, du2, ipiv

pytensor/tensor/slinalg.py

@@ -905,7 +905,7 @@ class Solve(SolveBase):
 
     def __init__(self, *, assume_a="gen", **kwargs):
         # Triangular and diagonal are handled outside of Solve
-        valid_options = ["gen", "sym", "her", "pos", "banded"]
+        valid_options = ["gen", "sym", "her", "pos", "banded", "tridiagonal"]
 
         assume_a = assume_a.lower()
         # We use the old names as the different dispatches are more likely to support them
@@ -922,7 +922,7 @@ def __init__(self, *, assume_a="gen", **kwargs):
                 f"Invalid assume_a: {assume_a}. It must be one of {valid_options} or {list(long_to_short.keys())}"
             )
 
-        if assume_a == "banded":
+        if assume_a in ("tridiagonal", "banded"):
             from scipy import __version__ as sp_version
 
             if tuple(map(int, sp_version.split(".")[:-1])) < (1, 15):
@@ -1043,14 +1043,14 @@ def solve(
             b_ndim=b_ndim,
         )
 
-    elif assume_a == "tridiagonal":
-        from pytensor.tensor._linalg.solve.tridiagonal import (
-            solve_tridiagonal_from_full_A_b,
-        )
-
-        return solve_tridiagonal_from_full_A_b(
-            a, b, b_ndim=b_ndim, transposed=transposed
-        )
+    # elif assume_a == "tridiagonal":
+    #     from pytensor.tensor._linalg.solve.tridiagonal import (
+    #         solve_tridiagonal_from_full_A_b,
+    #     )
+    #
+    #     return solve_tridiagonal_from_full_A_b(
+    #         a, b, b_ndim=b_ndim, transposed=transposed
+    #     )
 
     elif assume_a == "diagonal":
         a_diagonal = diagonal(a, axis1=-2, axis2=-1)