review comments: overwrite_a test + tri rewrite test

Author: asifzubair

pytensor/tensor/rewriting/linalg.py

@@ -1031,13 +1031,13 @@ def _find_triangular_op(var):
     """
     Inspects a variable to see if it's triangular.
 
-    Returns a tuple (is_lower, is_upper) if triangular, otherwise None.
+    Returns `True` if lower-triangular, `False` if upper-triangular, otherwise `None`.
     """
     # Case 1: Check for an explicit tag
     is_lower = getattr(var.tag, "lower_triangular", False)
     is_upper = getattr(var.tag, "upper_triangular", False)
     if is_lower or is_upper:
-        return (is_lower, is_upper)
+        return is_lower
 
     if not var.owner:
         return None
@@ -1047,7 +1047,7 @@ def _find_triangular_op(var):
 
     # Case 2: Check for direct creator Ops
     if isinstance(core_op, Cholesky):
-        return (core_op.lower, not core_op.lower)
+        return core_op.lower
 
     if isinstance(core_op, LU | LUFactor):
         if var.owner.outputs[1] == var:
@@ -1060,11 +1060,10 @@ def _find_triangular_op(var):
             return (False, True)
 
     # 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:
-    #         print('re-writing ... ')
-    #         return (True, False)
+    if isinstance(core_op, Tri):
+        k_node = var.owner.inputs[2]
+        if isinstance(k_node, Constant) and k_node.data == 0:
+            return True
 
     # Case 3: tril/triu patterns which are implemented as Mul
     if isinstance(core_op, Elemwise) and isinstance(core_op.scalar_op, ScalarMul):
@@ -1077,7 +1076,7 @@ def _find_triangular_op(var):
             if isinstance(other_inp.owner.op, Tri):
                 k_node = other_inp.owner.inputs[2]
                 if isinstance(k_node, Constant) and k_node.data == 0:
-                    return (True, False)  # It's tril
+                    return True  # It's tril
 
             # Check for triu pattern: Mul(x, Sub(1, Tri(k=-1)))
             sub_op = other_inp.owner.op
@@ -1095,7 +1094,7 @@ def _find_triangular_op(var):
                 if const_one is not None and tri_inp is not None:
                     k_node = tri_inp.owner.inputs[2]
                     if isinstance(k_node, Constant) and k_node.data == -1:
-                        return (False, True)  # It's triu
+                        return False  # It's triu
 
     return None
 
@@ -1111,13 +1110,11 @@ def rewrite_inv_to_triangular_solve(fgraph, node):
     """
 
     A = node.inputs[0]
-    triangular_info = _find_triangular_op(A)
-    if triangular_info is None:
+    is_lower = _find_triangular_op(A)
+    if is_lower is None:
         return None
 
-    is_lower, is_upper = triangular_info
-    if is_lower or is_upper:
-        new_op = TriangularInv(lower=is_lower)
-        new_inv = new_op(A)
-        copy_stack_trace(node.outputs[0], new_inv)
-        return [new_inv]
+    new_op = TriangularInv(lower=is_lower)
+    new_inv = new_op(A)
+    copy_stack_trace(node.outputs[0], new_inv)
+    return [new_inv]

pytensor/tensor/slinalg.py

@@ -1036,8 +1036,27 @@ def __init__(self, lower=True, on_error="raise", overwrite_a=False):
     def perform(self, node, inputs, outputs):
         (x,) = inputs
         (z,) = outputs
-        (dtrtri,) = get_lapack_funcs(("trtri",), (x,))
-        inv, info = dtrtri(x, lower=self.lower, overwrite_c=True)
+        (trtri,) = get_lapack_funcs(("trtri",), (x,))
+
+        # Check if we want to overwrite and if the input is C-contiguous
+        c_contiguous_input = self.overwrite_a and x.flags["C_CONTIGUOUS"]
+        if c_contiguous_input:
+            # Transpose C-contiguous to F-contiguous
+            x_in = x.T
+            lower_flag = not self.lower
+            overwrite_flag = True
+        else:
+            # Use original matrix and flags
+            x_in = x
+            lower_flag = self.lower
+            overwrite_flag = self.overwrite_a
+
+        # Call trtri with the potentially transposed input and correct flags
+        # Use overwrite_c (LAPACK flag for trtri) based on our logic
+        inv_maybe_transposed, info = trtri(
+            x_in, lower=lower_flag, overwrite_c=overwrite_flag
+        )
+
         if info != 0:
             if self.on_error == "nan":
                 z[0] = np.full_like(x, np.nan)
@@ -1048,7 +1067,7 @@ def perform(self, node, inputs, outputs):
                 raise ValueError(
                     f"illegal value in {-info}-th argument of internal trtri"
                 )
-        z[0] = inv
+        z[0] = inv_maybe_transposed.T if c_contiguous_input else inv_maybe_transposed
 
     def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
         """

tests/tensor/rewriting/test_linalg.py

@@ -14,7 +14,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor import swapaxes
-from pytensor.tensor.basic import tril, triu
+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
@@ -1085,9 +1085,9 @@ def _check_op_in_graph(fgraph, op_type, present=True):
 
 
 rewrite_cases = {
-    "tril": (
-        lambda x: tril(x),
-        lambda a: np.tril(a),
+    "tri": (
+        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),
     ),
     "triu": (
         lambda x: triu(x),

tests/tensor/test_slinalg.py

@@ -8,6 +8,7 @@
 import scipy
 from scipy import linalg as scipy_linalg
 
+import pytensor
 from pytensor import function, grad
 from pytensor import tensor as pt
 from pytensor.configdefaults import config
@@ -1238,43 +1239,57 @@ def _test_fn(x, case=2, mode="reduced"):
         )
 
 
-@pytest.mark.parametrize("lower", [True, False])
-def test_triangular_inv_op(lower):
-    """Tests the TriangularInv Op directly."""
-    x = matrix("x", dtype=config.floatX)
-    f = function([x], TriangularInv(lower=lower)(x))
-
-    if lower:
-        a = np.tril(np.random.rand(5, 5) + 0.1).astype(config.floatX)
-    else:
-        a = np.triu(np.random.rand(5, 5) + 0.1).astype(config.floatX)
-
-    a_inv = f(a)
-    expected_inv = np.linalg.inv(a)
-
-    # Clean the NumPy result before comparing.
-    if lower:
-        expected_inv = np.tril(expected_inv)
-    else:
-        expected_inv = np.triu(expected_inv)
-
-    # The inverse of a triangular matrix is also triangular.
-    # We should check the full matrix, not just a part of it.
-    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_nan_on_error():
+class TestTriangularInv:
     """
-    Tests the `on_error='nan'` functionality of the TriangularInv Op.
+    Tests for the `TriangularInv` `Op`.
     """
-    x = matrix("x", dtype=config.floatX)
-    f_nan = function([x], TriangularInv(on_error="nan")(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)
+    @pytest.mark.parametrize("lower", [True, False])
+    def test_triangular_inv_op(self, lower):
+        """Tests the TriangularInv Op directly."""
+        x = matrix("x", dtype=config.floatX)
+        f = function([x], TriangularInv(lower=lower)(x))
+
+        if lower:
+            a = np.tril(np.random.rand(5, 5) + 0.1).astype(config.floatX)
+        else:
+            a = np.triu(np.random.rand(5, 5) + 0.1).astype(config.floatX)
+
+        a_inv = f(a)
+        expected_inv = np.linalg.inv(a)
+
+        # The inverse of a triangular matrix is also triangular.
+        # We should check the full matrix, not just a part of it.
+        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_nan_on_error(self):
+        """
+        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))
+
+        # 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)
+
+        res = f_nan(a_singular)
+        assert np.all(np.isnan(res))
+
+    @pytest.mark.parametrize("overwrite_a", [True, False])
+    def test_triangular_inv_op_inplace(self, overwrite_a):
+        """Tests the TriangularInv Op directly."""
+        x = matrix("x", dtype=config.floatX)
+        f = function(
+            [pytensor.In(x, mutable=overwrite_a)],
+            TriangularInv(overwrite_a=overwrite_a)(x),
+            accept_inplace=True,
+        )
+
+        a = np.tril(np.random.rand(5, 5) + 0.1).astype(config.floatX)
+        a_copy = a.copy()
+        f(a)
 
-    res = f_nan(a_singular)
-    assert np.all(np.isnan(res))
+        assert overwrite_a == (not np.allclose(a, a_copy))