Use lapack instead of scipy_linalg.solve_triangular

Author: Fyrebright

pytensor/tensor/slinalg.py

@@ -8,7 +8,7 @@
 import scipy.linalg as scipy_linalg
 from numpy.exceptions import ComplexWarning
 from scipy.linalg import get_lapack_funcs
-from scipy.linalg._misc import LinAlgWarning
+from scipy.linalg._misc import LinAlgError, LinAlgWarning
 
 import pytensor
 from pytensor import ifelse
@@ -897,15 +897,51 @@ def __init__(self, *, unit_diagonal=False, **kwargs):
 
     def perform(self, node, inputs, outputs):
         A, b = inputs
-        outputs[0][0] = scipy_linalg.solve_triangular(
-            A,
-            b,
-            lower=self.lower,
-            trans=0,
-            unit_diagonal=self.unit_diagonal,
-            check_finite=self.check_finite,
-            overwrite_b=self.overwrite_b,
-        )
+
+        if self.check_finite and not (np.isfinite(A).all() and np.isfinite(b).all()):
+            raise ValueError("array must not contain infs or NaNs")
+
+        if len(A.shape) != 2 or A.shape[0] != A.shape[1]:
+            raise ValueError("expected square matrix")
+
+        if A.shape[0] != b.shape[0]:
+            raise ValueError(f"shapes of a {A.shape} and b {b.shape} are incompatible")
+
+        (trtrs,) = get_lapack_funcs(("trtrs",), (A, b))
+
+        # Quick return for empty arrays
+        if b.size == 0:
+            outputs[0][0] = np.empty_like(b, dtype=trtrs.dtype)
+            return
+
+        if A.flags["F_CONTIGUOUS"]:
+            x, info = trtrs(
+                A,
+                b,
+                overwrite_b=self.overwrite_b,
+                lower=self.lower,
+                trans=0,
+                unitdiag=self.unit_diagonal,
+            )
+        else:
+            # transposed system is solved since trtrs expects Fortran ordering
+            x, info = trtrs(
+                A.T,
+                b,
+                overwrite_b=self.overwrite_b,
+                lower=not self.lower,
+                trans=1,
+                unitdiag=self.unit_diagonal,
+            )
+
+        if info > 0:
+            raise LinAlgError(
+                f"singular matrix: resolution failed at diagonal {info-1}"
+            )
+        elif info < 0:
+            raise ValueError(f"illegal value in {-info}-th argument of internal trtrs")
+
+        outputs[0][0] = x
 
     def L_op(self, inputs, outputs, output_gradients):
         res = super().L_op(inputs, outputs, output_gradients)

tests/tensor/test_slinalg.py

@@ -513,6 +513,31 @@ def solve_op(A, b):
 
         utt.verify_grad(solve_op, [A_val, b_val], 3, rng, eps=eps)
 
+    def test_solve_triangular_empty(self):
+        rng = np.random.default_rng(utt.fetch_seed())
+        A = pt.tensor("A", shape=(5, 5))
+        b = pt.tensor("b", shape=(5, 0))
+
+        A_val = rng.random((5, 5)).astype(config.floatX)
+        b_empty = np.empty([5, 0], dtype=config.floatX)
+
+        A_func = functools.partial(self.A_func, lower=True, unit_diagonal=True)
+
+        x = solve_triangular(
+            A_func(A),
+            b,
+            lower=True,
+            trans=0,
+            unit_diagonal=True,
+            b_ndim=len((5, 0)),
+        )
+
+        f = function([A, b], x)
+
+        res = f(A_val, b_empty)
+        assert res.size == 0
+        assert res.dtype == config.floatX
+
 
 class TestCholeskySolve(utt.InferShapeTester):
     def setup_method(self):