Expand test coverage for Solve and SolveTriangular

Author: jessegrabowski

tests/link/jax/test_slinalg.py

@@ -5,6 +5,7 @@
 import pytest
 
 import pytensor.tensor as pt
+import tests.unittest_tools as utt
 from pytensor.configdefaults import config
 from pytensor.tensor import nlinalg as pt_nlinalg
 from pytensor.tensor import slinalg as pt_slinalg
@@ -103,28 +104,41 @@ def test_jax_basic():
     )
 
 
-@pytest.mark.parametrize("check_finite", [False, True])
-@pytest.mark.parametrize("lower", [False, True])
-@pytest.mark.parametrize("trans", [0, 1, 2])
-def test_jax_SolveTriangular(trans, lower, check_finite):
-    x = matrix("x")
-    b = vector("b")
+def test_jax_solve():
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("B", shape=(5, 5))
+
+    out = pt_slinalg.solve(A, b, lower=False, transposed=False)
+
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
 
-    out = pt_slinalg.solve_triangular(
-        x,
-        b,
-        trans=trans,
-        lower=lower,
-        check_finite=check_finite,
-    )
     compare_jax_and_py(
-        [x, b],
+        [A, b],
         [out],
-        [
-            np.eye(10).astype(config.floatX),
-            np.arange(10).astype(config.floatX),
-        ],
+        [A_val, b_val],
+    )
+
+
+def test_jax_SolveTriangular():
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("B", shape=(5, 5))
+
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
+
+    out = pt_slinalg.solve_triangular(
+        A,
+        b,
+        trans=0,
+        lower=True,
+        unit_diagonal=False,
     )
+    compare_jax_and_py([A, b], [out], [A_val, b_val])
 
 
 def test_jax_block_diag():

tests/link/numba/test_slinalg.py

@@ -5,7 +5,6 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_allclose
-from scipy import linalg as scipy_linalg
 
 import pytensor
 import pytensor.tensor as pt
@@ -26,9 +25,9 @@ def transpose_func(x, trans):
     if trans == 0:
         return x
     if trans == 1:
-        return x.conj().T
-    if trans == 2:
         return x.T
+    if trans == 2:
+        return x.conj().T
 
 
 @pytest.mark.parametrize(
@@ -59,18 +58,18 @@ def test_solve_triangular(b_shape: tuple[int], lower, trans, unit_diag, is_compl
 
     def A_func(x):
         x = x @ x.conj().T
-        x_tri = scipy_linalg.cholesky(x, lower=lower).astype(dtype)
+        x_tri = pt.linalg.cholesky(x, lower=lower).astype(dtype)
 
         if unit_diag:
-            x_tri[np.diag_indices_from(x_tri)] = 1.0
+            x_tri = pt.fill_diagonal(x_tri, 1.0)
 
-        return x_tri.astype(dtype)
+        return x_tri
 
     solve_op = partial(
         pt.linalg.solve_triangular, lower=lower, trans=trans, unit_diagonal=unit_diag
     )
 
-    X = solve_op(A, b)
+    X = solve_op(A_func(A), b)
     f = pytensor.function([A, b], X, mode="NUMBA")
 
     A_val = np.random.normal(size=(5, 5))
@@ -80,20 +79,20 @@ def A_func(x):
         A_val = A_val + np.random.normal(size=(5, 5)) * 1j
         b_val = b_val + np.random.normal(size=b_shape) * 1j
 
-    X_np = f(A_func(A_val), b_val)
-
-    test_input = transpose_func(A_func(A_val), trans)
-
-    ATOL = 1e-8 if floatX.endswith("64") else 1e-4
-    RTOL = 1e-8 if floatX.endswith("64") else 1e-4
-
-    np.testing.assert_allclose(test_input @ X_np, b_val, atol=ATOL, rtol=RTOL)
+    X_np = f(A_val.copy(), b_val.copy())
+    A_val_transformed = transpose_func(A_func(A_val), trans).eval()
+    np.testing.assert_allclose(
+        A_val_transformed @ X_np,
+        b_val,
+        atol=1e-8 if floatX.endswith("64") else 1e-4,
+        rtol=1e-8 if floatX.endswith("64") else 1e-4,
+    )
 
     compiled_fgraph = f.maker.fgraph
     compare_numba_and_py(
         compiled_fgraph.inputs,
         compiled_fgraph.outputs,
-        [A_func(A_val), b_val],
+        [A_val, b_val],
     )
 
 
@@ -145,7 +144,6 @@ def test_solve_triangular_overwrite_b_correct(overwrite_b):
     b_test_nb = b_test_py.copy(order="F")
 
     op = SolveTriangular(
-        trans=0,
         unit_diagonal=False,
         lower=False,
         check_finite=True,