Incorporate feedback

Author: jessegrabowski

pytensor/link/numba/dispatch/slinalg.py

@@ -126,17 +126,21 @@ def impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b):
 
         B_is_1d = B.ndim == 1
 
-        A_copy = _copy_to_fortran_order(A)
-
-        # This list is exhaustive, but numba freaks out if we include a final else clause
-        if not overwrite_b and not B_is_1d:
-            B_copy = _copy_to_fortran_order(B)
-        elif overwrite_b and not B_is_1d:
-            B_copy = np.asfortranarray(B)
-        elif not overwrite_b and B_is_1d:
-            B_copy = np.copy(np.expand_dims(B, -1))
-        elif overwrite_b and B_is_1d:
-            B_copy = np.expand_dims(B, -1)
+        # Despite the name, this won't copy if A is already fortran contiguous
+        A_copy = np.asfortranarray(A)
+
+        if overwrite_b:
+            if B_is_1d:
+                B_copy = np.expand_dims(B, -1)
+            else:
+                # Same here, this *will* allow inplace destruction of B, but only if it is already fortran contiguous.
+                # Otherwise, there's no way to get around the need to copy the data before going into TRTRS
+                B_copy = np.asfortranarray(B)
+        else:
+            if B_is_1d:
+                B_copy = np.copy(np.expand_dims(B, -1))
+            else:
+                B_copy = _copy_to_fortran_order(B)
 
         NRHS = 1 if B_is_1d else int(B_copy.shape[-1])
 

tests/link/numba/test_slinalg.py

@@ -141,8 +141,8 @@ def test_solve_triangular_overwrite_b_correct(overwrite_b):
     b_test_py = np.asfortranarray(rng.normal(size=(3, 2)))
 
     # .T.copy().T creates an f-contiguous copy of an f-contiguous array (otherwise the copy is c-contiguous)
-    a_test_nb = a_test_py.T.copy().T
-    b_test_nb = b_test_py.T.copy().T
+    a_test_nb = a_test_py.copy(order="F")
+    b_test_nb = b_test_py.copy(order="F")
 
     op = SolveTriangular(
         trans=0,