Incorporate feedback

Author: jessegrabowski

pytensor/link/numba/dispatch/linalg/dot/banded.py

@@ -59,6 +59,7 @@ def impl(A: np.ndarray, x: np.ndarray, kl: int, ku: int) -> np.ndarray:
         m, n = A.shape
 
         A_banded = A_to_banded(A, kl=kl, ku=ku)
+        stride = x.strides[0] // x.itemsize
 
         TRANS = val_to_int_ptr(ord("N"))
         M = val_to_int_ptr(m)
@@ -69,7 +70,8 @@ def impl(A: np.ndarray, x: np.ndarray, kl: int, ku: int) -> np.ndarray:
         KU = val_to_int_ptr(ku)
 
         ALPHA = np.array(1.0, dtype=dtype)
-        INCX = val_to_int_ptr(x.strides[0] // x.itemsize)
+
+        INCX = val_to_int_ptr(stride)
         BETA = np.array(0.0, dtype=dtype)
         Y = np.empty(m, dtype=dtype)
         INCY = val_to_int_ptr(1)

pytensor/tensor/slinalg.py

@@ -1722,9 +1722,7 @@ def L_op(self, inputs, outputs, output_grads) -> list[Variable]:
         (G_bar,) = output_grads
 
         A_bar = pt.outer(G_bar, x.T)
-        x_bar = banded_dot(
-            A.T, G_bar, lower_diags=self.lower_diags, upper_diags=self.upper_diags
-        )
+        x_bar = self(A.T, G_bar)
 
         return [A_bar, x_bar]
 

tests/link/numba/test_slinalg.py

@@ -724,7 +724,7 @@ def test_lu_solve(b_func, b_shape: tuple[int, ...], trans: bool, overwrite_b: bo
     np.testing.assert_allclose(b_val_not_contig, b_val)
 
 
-@pytest.mark.parametrize("stride", [1, 2, -1], ids=lambda x: f"stride={x}")
+@pytest.mark.parametrize("stride", [1, 2, -1, -2], ids=lambda x: f"stride={x}")
 def test_banded_dot(stride):
     rng = np.random.default_rng()
 
@@ -743,19 +743,6 @@ def test_banded_dot(stride):
         [A, x],
         output,
         test_inputs=[A_val, x_val],
-        inplace=True,
-        numba_mode=numba_inplace_mode,
-        eval_obj_mode=False,
-    )
-
-    # Test non-contiguous x input
-    x_val = rng.normal(size=(20,))[::2]
-
-    compare_numba_and_py(
-        [A, x],
-        output,
-        test_inputs=[A_val, x_val],
-        inplace=True,
         numba_mode=numba_inplace_mode,
         eval_obj_mode=False,
     )