Decompose Tridiagonal Solve into core steps

Author: ricardoV94

pytensor/link/numba/dispatch/linalg/solve/tridiagonal.py

@@ -4,7 +4,6 @@
 from numba.core.extending import overload
 from numba.np.linalg import ensure_lapack
 from numpy import ndarray
-from scipy import linalg
 
 from pytensor.link.numba.dispatch.basic import numba_njit
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
@@ -13,11 +12,9 @@
     int_ptr_to_val,
     val_to_int_ptr,
 )
-from pytensor.link.numba.dispatch.linalg.solve.utils import _solve_check_input_shapes
 from pytensor.link.numba.dispatch.linalg.utils import (
     _check_scipy_linalg_matrix,
     _copy_to_fortran_order_even_if_1d,
-    _solve_check,
     _trans_char_to_int,
 )
 
@@ -227,72 +224,42 @@ def impl(
 
 
 def _solve_tridiagonal(
-    a: ndarray,
-    b: ndarray,
-    lower: bool,
-    overwrite_a: bool,
+    dl: ndarray,
+    d: ndarray,
+    ul: ndarray,
+    B: ndarray,
     overwrite_b: bool,
-    check_finite: bool,
-    transposed: bool,
 ):
     """
-    Solve a positive-definite linear system using the Cholesky decomposition.
+    Solve a tridiagonal linear system.
     """
-    return linalg.solve(
-        a=a,
-        b=b,
-        lower=lower,
-        overwrite_a=overwrite_a,
-        overwrite_b=overwrite_b,
-        check_finite=check_finite,
-        transposed=transposed,
-        assume_a="tridiagonal",
-    )
+    return
 
 
 @overload(_solve_tridiagonal)
-def _tridiagonal_solve_impl(
-    A: ndarray,
+def _solve_tridiagonal_impl(
+    dl: ndarray,
+    d: ndarray,
+    du: ndarray,
     B: ndarray,
-    lower: bool,
-    overwrite_a: bool,
     overwrite_b: bool,
-    check_finite: bool,
-    transposed: bool,
-) -> Callable[[ndarray, ndarray, bool, bool, bool, bool, bool], ndarray]:
+) -> Callable[[ndarray, ndarray, ndarray, ndarray, bool], ndarray]:
     ensure_lapack()
-    _check_scipy_linalg_matrix(A, "solve")
+    _check_scipy_linalg_matrix(dl, "solve_")
+    _check_scipy_linalg_matrix(dl, "solve")
+    _check_scipy_linalg_matrix(dl, "solve")
     _check_scipy_linalg_matrix(B, "solve")
 
     def impl(
-        A: ndarray,
+        dl: ndarray,
+        d: ndarray,
+        du: ndarray,
         B: ndarray,
-        lower: bool,
-        overwrite_a: bool,
         overwrite_b: bool,
-        check_finite: bool,
-        transposed: bool,
     ) -> ndarray:
-        n = np.int32(A.shape[-1])
-        _solve_check_input_shapes(A, B)
-        norm = "1"
-
-        if transposed:
-            A = A.T
-        dl, d, du = np.diag(A, -1), np.diag(A, 0), np.diag(A, 1)
-
-        anorm = tridiagonal_norm(du, d, dl)
-
-        dl, d, du, du2, IPIV, INFO = _gttrf(dl, d, du)
-        _solve_check(n, INFO)
-
-        X, INFO = _gttrs(
-            dl, d, du, du2, IPIV, B, trans=transposed, overwrite_b=overwrite_b
-        )
-        _solve_check(n, INFO)
+        dl, d, du, du2, IPIV, _ = _gttrf(dl, d, du)
 
-        RCOND, INFO = _gtcon(dl, d, du, du2, IPIV, anorm, norm)
-        _solve_check(n, INFO, True, RCOND)
+        X, _ = _gttrs(dl, d, du, du2, IPIV, B, trans=0, overwrite_b=overwrite_b)
 
         return X
 

pytensor/link/numba/dispatch/slinalg.py

@@ -17,6 +17,7 @@
 from pytensor.link.numba.dispatch.linalg.solve.symmetric import _solve_symmetric
 from pytensor.link.numba.dispatch.linalg.solve.triangular import _solve_triangular
 from pytensor.link.numba.dispatch.linalg.solve.tridiagonal import _solve_tridiagonal
+from pytensor.tensor._linalg.solve.tridiagonal import SolveTridiagonal
 from pytensor.tensor.slinalg import (
     LU,
     BlockDiagonal,
@@ -215,8 +216,6 @@ def numba_funcify_Solve(op, node, **kwargs):
         solve_fn = _solve_symmetric
     elif assume_a == "pos":
         solve_fn = _solve_psd
-    elif assume_a == "tridiagonal":
-        solve_fn = _solve_tridiagonal
     else:
         warnings.warn(
             f"Numba assume_a={assume_a} not implemented. Falling back to general solve.\n"
@@ -284,6 +283,32 @@ def solve_triangular(a, b):
     return solve_triangular
 
 
+@numba_funcify.register(SolveTridiagonal)
+def numba_funcify_SolveTridiagonal(op, node, **kwargs):
+    overwrite_b = op.overwrite_b
+    b_ndim = op.b_ndim
+
+    dtype = node.inputs[0].dtype
+    if dtype in complex_dtypes:
+        raise NotImplementedError(
+            _COMPLEX_DTYPE_NOT_SUPPORTED_MSG.format(op="Solve Triangular")
+        )
+
+    @numba_njit
+    def solve_triangular(dl, d, du, b):
+        res = _solve_tridiagonal(
+            dl,
+            d,
+            du,
+            b,
+            overwrite_b=overwrite_b,
+        )
+
+        return res
+
+    return solve_triangular
+
+
 @numba_funcify.register(CholeskySolve)
 def numba_funcify_CholeskySolve(op, node, **kwargs):
     lower = op.lower

pytensor/scan/rewriting.py

@@ -2561,26 +2561,24 @@ def scan_push_out_dot1(fgraph, node):
     position=1,
 )
 
-
 scan_seqopt1.register(
     "scan_push_out_non_seq",
     in2out(scan_push_out_non_seq, ignore_newtrees=True),
     "scan_pushout_nonseqs_ops",  # For backcompat: so it can be tagged with old name
     "fast_run",
     "scan",
     "scan_pushout",
-    position=2,
+    position=3,
 )
 
-
 scan_seqopt1.register(
     "scan_push_out_seq",
     in2out(scan_push_out_seq, ignore_newtrees=True),
     "scan_pushout_seqs_ops",  # For backcompat: so it can be tagged with old name
     "fast_run",
     "scan",
     "scan_pushout",
-    position=3,
+    position=4,
 )
 
 
@@ -2592,7 +2590,7 @@ def scan_push_out_dot1(fgraph, node):
     "more_mem",
     "scan",
     "scan_pushout",
-    position=4,
+    position=5,
 )
 
 
@@ -2605,7 +2603,7 @@ def scan_push_out_dot1(fgraph, node):
     "more_mem",
     "scan",
     "scan_pushout",
-    position=5,
+    position=6,
 )
 
 scan_eqopt2.register(

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.rewriting