Improve solve test

Author: jessegrabowski

tests/tensor/test_slinalg.py

@@ -238,37 +238,46 @@ def test_infer_shape(self, b_shape):
     "b_size", [(5, 1), (5, 5), (5,)], ids=["b_col_vec", "b_matrix", "b_vec"]
 )
 @pytest.mark.parametrize("assume_a", ["gen", "sym", "pos"], ids=str)
-@pytest.mark.parametrize("transposed", [True, False], ids=["trans", "no_trans"])
-def test_solve_correctness(b_size: tuple[int], assume_a: str, transposed: bool):
+def test_solve_correctness(b_size: tuple[int], assume_a: str):
     rng = np.random.default_rng(utt.fetch_seed())
     A = pt.tensor("A", shape=(5, 5))
     b = pt.tensor("b", shape=b_size)
 
-    solve_op = functools.partial(
-        solve, assume_a=assume_a, transposed=transposed, b_ndim=len(b_size)
-    )
-    y = solve_op(
-        A,
-        b,
-    )
-    solve_func = pytensor.function([A, b], y)
-
-    b_val = rng.normal(size=b_size).astype(config.floatX)
     A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=b_size).astype(config.floatX)
+
+    solve_op = functools.partial(solve, assume_a=assume_a, b_ndim=len(b_size))
 
-    if assume_a == "sym":
-        A_val = (A_val + A_val.T) / 2
-    elif assume_a == "pos":
-        A_val = A_val @ A_val.T
+    def A_func(x):
+        if assume_a == "pos":
+            return x @ x.T
+        elif assume_a == "sym":
+            return (x + x.T) / 2
+        else:
+            return x
+
+    solve_input_val = A_func(A_val)
+
+    y = solve_op(A_func(A), b)
+    solve_func = pytensor.function([A, b], y)
+    X_np = solve_func(A_val.copy(), b_val.copy())
 
     np.testing.assert_allclose(
-        scipy.linalg.solve(A_val, b_val, assume_a=assume_a, transposed=transposed),
-        solve_func(A_val, b_val),
+        scipy.linalg.solve(solve_input_val, b_val, assume_a=assume_a),
+        X_np,
     )
 
+    np.testing.assert_allclose(A_func(A_val) @ X_np, b_val, atol=1e-6)
+
     eps = 2e-8 if config.floatX == "float64" else None
 
-    utt.verify_grad(solve_op, [A_val, b_val], 3, rng, eps=eps)
+    # To correctly check the gradients, we need to include a transformation from the space of unconstrained matrices
+    # (A) to a valid input matrix for the given solver. This is done by the A_func function. If this isn't included,
+    # the random perturbations used by verify_grad will result in invalid input matrices, and
+    # LAPACK will silently do the wrong thing, making the gradients wrong
+    utt.verify_grad(
+        lambda A, b: solve_op(A_func(A), b), [A_val, b_val], 3, rng, eps=eps
+    )
 
 
 class TestSolveTriangular(utt.InferShapeTester):