Speedup solve tridiagonal

Author: ricardoV94

pytensor/tensor/slinalg.py

@@ -937,15 +937,24 @@ def __init__(self, *, assume_a="gen", **kwargs):
 
     def perform(self, node, inputs, outputs):
         a, b = inputs
-        outputs[0][0] = scipy_linalg.solve(
-            a=a,
-            b=b,
-            lower=self.lower,
-            check_finite=self.check_finite,
-            assume_a=self.assume_a,
-            overwrite_a=self.overwrite_a,
-            overwrite_b=self.overwrite_b,
-        )
+        if self.assume_a == "tridiagonal":
+            [dl, d, du] = (a.diagonal(offset=o) for o in (-1, 0, 1))
+            _gttrf, _gttrs = scipy_linalg.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)
+            outputs[0][0] = x
+        else:
+            outputs[0][0] = scipy_linalg.solve(
+                a=a,
+                b=b,
+                lower=self.lower,
+                check_finite=self.check_finite,
+                assume_a=self.assume_a,
+                overwrite_a=self.overwrite_a,
+                overwrite_b=self.overwrite_b,
+            )
 
     def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
         if not allowed_inplace_inputs: