Adjust numba test

Author: jessegrabowski

tests/link/numba/test_slinalg.py

@@ -1,5 +1,6 @@
 import re
 from typing import Literal
+from typing import cast as typing_cast
 
 import numpy as np
 import pytest
@@ -724,25 +725,36 @@ 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, -2], ids=lambda x: f"stride={x}")
-def test_banded_dot(stride):
+def test_banded_dot():
     rng = np.random.default_rng()
 
+    A = pt.tensor("A", shape=(10, 10), dtype=config.floatX)
     A_val = _make_banded_A(rng.normal(size=(10, 10)), kl=1, ku=1).astype(config.floatX)
 
-    x_shape = (10 * abs(stride),)
-    x_val = rng.normal(size=x_shape).astype(config.floatX)
-    x_val = x_val[::stride]
-
-    A = pt.tensor("A", shape=A_val.shape, dtype=A_val.dtype)
-    x = pt.tensor("x", shape=x_val.shape, dtype=x_val.dtype)
+    x = pt.tensor("x", shape=(10,), dtype=config.floatX)
+    x_val = rng.normal(size=(10,)).astype(config.floatX)
 
     output = banded_dot(A, x, upper_diags=1, lower_diags=1)
 
-    compare_numba_and_py(
+    fn, _ = compare_numba_and_py(
         [A, x],
         output,
         test_inputs=[A_val, x_val],
         numba_mode=numba_inplace_mode,
         eval_obj_mode=False,
     )
+
+    for stride in [2, -1, -2]:
+        x_shape = (10 * abs(stride),)
+        x_val = rng.normal(size=x_shape).astype(config.floatX)
+        x_val = x_val[::stride]
+
+        nb_output = typing_cast(np.ndarray, fn(A_val, x_val))
+        expected = A_val @ x_val
+
+        np.testing.assert_allclose(
+            nb_output,
+            expected,
+            strict=True,
+            err_msg=f"Test failed for stride = {stride}",
+        )