improve test coverage, fix return types

Author: asifzubair

pytensor/tensor/rewriting/linalg.py

@@ -1051,15 +1051,14 @@ def _find_triangular_op(var):
 
     if isinstance(core_op, LU | LUFactor):
         if var.owner.outputs[1] == var:
-            return (True, False)
+            return True
         if var.owner.outputs[2] == var:
-            return (False, True)
+            return False
 
     if isinstance(core_op, QR):
         if var.owner.outputs[1] == var:
-            return (False, True)
+            return False
 
-    # pt.tri will get constant folded so no point re-writing ?
     if isinstance(core_op, Tri):
         k_node = var.owner.inputs[2]
         if isinstance(k_node, Constant) and k_node.data == 0:

tests/tensor/rewriting/test_linalg.py

@@ -14,7 +14,6 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor import swapaxes
-from pytensor.tensor.basic import triu
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.math import dot, matmul
@@ -1089,8 +1088,12 @@ def _check_op_in_graph(fgraph, op_type, present=True):
         lambda x: pt.tri(x.shape[0], x.shape[1], k=0, dtype=x.dtype),
         lambda a: np.tri(N=a.shape[0], M=a.shape[1], k=0, dtype=a.dtype),
     ),
+    "tril": (
+        lambda x: pt.tril(x),
+        lambda a: np.tril(a),
+    ),
     "triu": (
-        lambda x: triu(x),
+        lambda x: pt.triu(x),
         lambda a: np.triu(a),
     ),
     "cholesky": (

tests/tensor/test_slinalg.py

@@ -1263,10 +1263,26 @@ def test_triangular_inv_op(self, lower):
         atol = rtol = 1e-8 if config.floatX.endswith("64") else 1e-4
         np.testing.assert_allclose(a_inv, expected_inv, rtol=rtol, atol=atol)
 
+    def test_triangular_inv_op_bad_on_error(self):
+        """Tests that a bad `on_error` value raises a ValueError."""
+        with pytest.raises(ValueError, match="on_error must be one of"):
+            TriangularInv(on_error="foo")
+
+    def test_triangular_inv_op_raise_on_error(self):
+        """Tests the default `on_error='raise'` functionality."""
+        x = matrix("x", dtype=config.floatX)
+        f_raise = function([x], TriangularInv()(x))
+
+        # Create a singular triangular matrix (zero on the diagonal)
+        a_singular = np.tril(np.random.rand(5, 5))
+        a_singular[2, 2] = 0
+        a_singular = a_singular.astype(config.floatX)
+
+        with pytest.raises(np.linalg.LinAlgError, match="Singular matrix"):
+            f_raise(a_singular)
+
     def test_triangular_inv_op_nan_on_error(self):
-        """
-        Tests the `on_error='nan'` functionality of the TriangularInv Op.
-        """
+        """Tests the `on_error='nan'` functionality of the TriangularInv Op."""
         x = matrix("x", dtype=config.floatX)
         f_nan = function([x], TriangularInv(on_error="nan")(x))