Fix numba dispatch of Det and SlogDet returning non-arrays

Author: ricardoV94

pytensor/link/numba/dispatch/nlinalg.py

@@ -52,7 +52,7 @@ def numba_funcify_Det(op, node, **kwargs):
 
     @numba_basic.numba_njit(inline="always")
     def det(x):
-        return numba_basic.direct_cast(np.linalg.det(inputs_cast(x)), out_dtype)
+        return np.array(np.linalg.det(inputs_cast(x))).astype(out_dtype)
 
     return det
 
@@ -68,8 +68,8 @@ def numba_funcify_SLogDet(op, node, **kwargs):
     def slogdet(x):
         sign, det = np.linalg.slogdet(inputs_cast(x))
         return (
-            numba_basic.direct_cast(sign, out_dtype_1),
-            numba_basic.direct_cast(det, out_dtype_2),
+            np.array(sign).astype(out_dtype_1),
+            np.array(det).astype(out_dtype_2),
         )
 
     return slogdet

tests/link/numba/test_basic.py

@@ -260,9 +260,12 @@ def compare_numba_and_py(
     if assert_fn is None:
 
         def assert_fn(x, y):
-            return np.testing.assert_allclose(x, y, rtol=1e-4) and compare_shape_dtype(
-                x, y
-            )
+            np.testing.assert_allclose(x, y, rtol=1e-4, strict=True)
+            # Make sure we don't have one input be a np.ndarray while the other is not
+            if isinstance(x, np.ndarray):
+                assert isinstance(y, np.ndarray), "y is not a NumPy array, but x is"
+            else:
+                assert not isinstance(y, np.ndarray), "y is a NumPy array, but x is not"
 
     if any(
         inp.owner is not None

tests/link/numba/test_nlinalg.py

@@ -11,68 +11,18 @@
 rng = np.random.default_rng(42849)
 
 
-@pytest.mark.parametrize(
-    "x, exc",
-    [
-        (
-            (
-                pt.dmatrix(),
-                (lambda x: x.T.dot(x))(rng.random(size=(3, 3)).astype("float64")),
-            ),
-            None,
-        ),
-        (
-            (
-                pt.lmatrix(),
-                (lambda x: x.T.dot(x))(rng.poisson(size=(3, 3)).astype("int64")),
-            ),
-            None,
-        ),
-    ],
-)
-def test_Det(x, exc):
-    x, test_x = x
-    g = nlinalg.Det()(x)
+@pytest.mark.parametrize("dtype", ("float64", "int64"))
+@pytest.mark.parametrize("op", (nlinalg.Det(), nlinalg.SLogDet()))
+def test_Det_SLogDet(op, dtype):
+    x = pt.matrix(dtype=dtype)
 
-    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
-    with cm:
-        compare_numba_and_py(
-            [x],
-            g,
-            [test_x],
-        )
+    rng = np.random.default_rng([50, sum(map(ord, dtype))])
+    x_ = rng.random(size=(3, 3)).astype(dtype)
+    test_x = x_.T.dot(x_)
 
+    g = op(x)
 
-@pytest.mark.parametrize(
-    "x, exc",
-    [
-        (
-            (
-                pt.dmatrix(),
-                (lambda x: x.T.dot(x))(rng.random(size=(3, 3)).astype("float64")),
-            ),
-            None,
-        ),
-        (
-            (
-                pt.lmatrix(),
-                (lambda x: x.T.dot(x))(rng.poisson(size=(3, 3)).astype("int64")),
-            ),
-            None,
-        ),
-    ],
-)
-def test_SLogDet(x, exc):
-    x, test_x = x
-    g = nlinalg.SLogDet()(x)
-
-    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
-    with cm:
-        compare_numba_and_py(
-            [x],
-            g,
-            [test_x],
-        )
+    compare_numba_and_py([x], g, [test_x])
 
 
 # We were seeing some weird results in CI where the following two almost