Avoid copying C-contiguous arrays in solve methods that only work with a triangular half of the matrix

Author: ricardoV94

pytensor/link/numba/dispatch/slinalg.py

@@ -136,7 +136,15 @@ def impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b):
         B_is_1d = B.ndim == 1
 
         # This will only copy if A is not already fortran contiguous
-        A_f = np.asfortranarray(A)
+        if A.flags.f_contiguous or (A.flags.c_contiguous and trans in (0, 1)):
+            A_f = A
+            if A.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                # Is this valid for complex matrices that were .conj().mT by PyTensor?
+                lower = not lower
+                trans = 1 - trans
+        else:
+            A_f = np.asfortranarray(A)
 
         if overwrite_b and B.flags.f_contiguous:
             B_copy = B
@@ -682,8 +690,11 @@ def impl(
         _LDA, _N = np.int32(A.shape[-2:])  # type: ignore
         _solve_check_input_shapes(A, B)
 
-        if overwrite_a and A.flags.f_contiguous:
+        if overwrite_a and (A.flags.f_contiguous or A.flags.c_contiguous):
             A_copy = A
+            if A.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                lower = not lower
         else:
             A_copy = _copy_to_fortran_order(A)
 
@@ -905,8 +916,11 @@ def impl(
 
         _N = np.int32(A.shape[-1])
 
-        if overwrite_a and A.flags.f_contiguous:
+        if overwrite_a and (A.flags.f_contiguous or A.flags.c_contiguous):
             A_copy = A
+            if A.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                lower = not lower
         else:
             A_copy = _copy_to_fortran_order(A)
 
@@ -1128,7 +1142,13 @@ def impl(C, B, lower=False, overwrite_b=False, check_finite=True):
         _solve_check_input_shapes(C, B)
 
         _N = np.int32(C.shape[-1])
-        C_f = np.asfortranarray(C)
+        if C.flags.f_contiguous or C.flags.c_contiguous:
+            C_f = C
+            if C.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                lower = not lower
+        else:
+            C_f = np.asfortranarray(C)
 
         if overwrite_b and B.flags.f_contiguous:
             B_copy = B

tests/link/numba/test_slinalg.py

@@ -88,7 +88,7 @@ class TestSolves:
         [(5, 1), (5, 5), (5,)],
         ids=["b_col_vec", "b_matrix", "b_vec"],
     )
-    @pytest.mark.parametrize("assume_a", ["gen", "sym", "pos"], ids=str)
+    @pytest.mark.parametrize("assume_a", ["gen", "sym", "pos"][1::2], ids=str)
     @pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower={x}")
     @pytest.mark.parametrize(
         "overwrite_a, overwrite_b",
@@ -169,7 +169,11 @@ def A_func(x):
         b_val_c_contig = np.copy(b_val, order="C")
         res_c_contig = f(A_val_c_contig, b_val_c_contig)
         np.testing.assert_allclose(res_c_contig, res)
-        np.testing.assert_allclose(A_val_c_contig, A_val)
+        # In the symmetric and positive definite cases,
+        # we can only destroy A C-contiguous arrays by inverting `lower` at runtime
+        assert np.allclose(A_val_c_contig, A_val) == (
+            not (overwrite_a and assume_a in ("sym", "pos"))
+        )
         # b vectors are always f_contiguous if also c_contiguous
         assert np.allclose(b_val_c_contig, b_val) == (
             not (overwrite_b and b_val_c_contig.flags.f_contiguous)
@@ -273,7 +277,7 @@ def A_func(x):
         b_val_c_contig = np.copy(b_val, order="C")
         res_c_contig = f(A_val_c_contig, b_val_c_contig)
         np.testing.assert_allclose(res_c_contig, res)
-        np.testing.assert_allclose(A_val_c_contig, A_val)
+        assert np.allclose(A_val_c_contig, A_val)
         # b c_contiguous vectors are also f_contiguous and destroyable
         assert np.allclose(b_val_c_contig, b_val) == (
             not (overwrite_b and b_val_c_contig.flags.f_contiguous)
@@ -359,7 +363,7 @@ def A_func(x):
         res_f_contig = f(A_val_f_contig, b_val_f_contig)
         np.testing.assert_allclose(res_f_contig, res)
         # cho_solve never destroys A
-        np.testing.assert_allclose(A_val, A_val_f_contig)
+        np.testing.assert_allclose(A_val == A_val_f_contig)
         # b Should always be destroyable
         assert (b_val == b_val_f_contig).all() == (not overwrite_b)
 
@@ -368,7 +372,7 @@ def A_func(x):
         b_val_c_contig = np.copy(b_val, order="C")
         res_c_contig = f(A_val_c_contig, b_val_c_contig)
         np.testing.assert_allclose(res_c_contig, res)
-        np.testing.assert_allclose(A_val_c_contig, A_val)
+        assert np.allclose(A_val_c_contig, A_val)
         # b c_contiguous vectors are also f_contiguous and destroyable
         assert np.allclose(b_val_c_contig, b_val) == (
             not (overwrite_b and b_val_c_contig.flags.f_contiguous)