Tests pass

Author: jessegrabowski

pytensor/link/numba/dispatch/slinalg.py

@@ -333,7 +333,7 @@ def block_diag(*arrs):
     return block_diag
 
 
-def _xlamch():
+def _xlamch(kind: str = "E"):
     """
     Placeholder for getting machine precision; used by linalg.solve. Not used by pytensor to numbify graphs.
     """
@@ -830,8 +830,10 @@ def impl(
         transposed: bool,
     ) -> np.ndarray:
         _solve_check_input_shapes(A, B)
+
         x, ipiv, info = _sysv(A, B, lower, overwrite_a, overwrite_b)
         _solve_check(A.shape[-1], info)
+
         rcond, info = _sycon(A, ipiv, _xlange(A, order="I"))
         _solve_check(A.shape[-1], info, True, rcond)
 
@@ -1024,8 +1026,10 @@ def impl(
         transposed: bool,
     ) -> np.ndarray:
         _solve_check_input_shapes(A, B)
+
         x, info = _posv(A, B, lower, overwrite_a, overwrite_b, check_finite, transposed)
         _solve_check(A.shape[-1], info)
+
         rcond, info = _pocon(x, _xlange(A))
         _solve_check(A.shape[-1], info=info, lamch=True, rcond=rcond)
 

tests/link/numba/test_slinalg.py

@@ -182,13 +182,12 @@ def test_numba_Cholesky_raise_on(on_error):
 
 
 @pytest.mark.parametrize("lower", [True, False], ids=["lower=True", "lower=False"])
-@pytest.mark.parametrize("trans", [True, False], ids=["trans=True", "trans=False"])
-def test_numba_Cholesky_grad(lower, trans):
+def test_numba_Cholesky_grad(lower):
     rng = np.random.default_rng(utt.fetch_seed())
     L = rng.normal(size=(5, 5)).astype(floatX)
     X = L @ L.T
 
-    chol_op = partial(pt.linalg.cholesky, lower=lower, trans=trans)
+    chol_op = partial(pt.linalg.cholesky, lower=lower)
     utt.verify_grad(chol_op, [X], mode="NUMBA")
 
 
@@ -359,10 +358,9 @@ def test_solve(b_shape: tuple[int], assume_a: Literal["gen", "sym", "pos"]):
 
     def A_func(x):
         if assume_a == "pos":
-            x = x.T @ x
+            x = x @ x.T
         elif assume_a == "sym":
-            x = (x.T + x) / 2
-
+            x = (x + x.T) / 2
         return x
 
     X = pt.linalg.solve(
@@ -376,13 +374,15 @@ def A_func(x):
     )
     op = f.maker.fgraph.outputs[0].owner.op
 
-    compare_numba_and_py(f.maker.fgraph, inputs=[A_func(A_val.copy()), b_val.copy()])
+    compare_numba_and_py(
+        ([A, b], [X]), inputs=[A_val.copy(), b_val.copy()], inplace=True
+    )
 
     # Calling this is destructive and will rewrite b_val to be the answer. Store copies of the inputs first.
     A_val_copy = A_val.copy()
     b_val_copy = b_val.copy()
 
-    X_np = f(A_func(A_val), b_val)
+    X_np = f(A_val, b_val)
 
     # overwrite_b is preferred when both inputs can be destroyed
     assert op.destroy_map == {0: [1]}