Respect transposed argument in JAX, improve jax sovle tests

Author: jessegrabowski

pytensor/link/jax/dispatch/slinalg.py

@@ -39,13 +39,26 @@ def cholesky(a, lower=lower):
 
 @jax_funcify.register(Solve)
 def jax_funcify_Solve(op, **kwargs):
-    if op.assume_a != "gen" and op.lower:
-        lower = True
-    else:
-        lower = False
+    assume_a = op.assume_a
+    lower = op.lower
+    check_finite = op.check_finite
+    overwrite_a = op.overwrite_a
+    overwrite_b = op.overwrite_b
+    transposed = op.transposed
+
+    def solve(a, b):
+        if transposed:
+            a = a.T
 
-    def solve(a, b, lower=lower):
-        return jax.scipy.linalg.solve(a, b, lower=lower)
+        return jax.scipy.linalg.solve(
+            a,
+            b,
+            assume_a=assume_a,
+            lower=lower,
+            check_finite=check_finite,
+            overwrite_a=overwrite_a,
+            overwrite_b=overwrite_b,
+        )
 
     return solve
 

tests/link/jax/test_slinalg.py

@@ -73,16 +73,6 @@ def test_jax_basic():
         ],
     )
 
-    out = pt_slinalg.solve(x, b)
-    compare_jax_and_py(
-        [x, b],
-        [out],
-        [
-            np.eye(10).astype(config.floatX),
-            np.arange(10).astype(config.floatX),
-        ],
-    )
-
     out = pt.diag(b)
     compare_jax_and_py([b], [out], [np.arange(10).astype(config.floatX)])
 
@@ -103,6 +93,45 @@ def test_jax_basic():
     )
 
 
+solve_test_cases = [
+    ("gen", False, False),
+    ("gen", False, True),
+    ("sym", False, False),
+    ("sym", True, False),
+    ("sym", True, True),
+    ("pos", False, False),
+    ("pos", True, False),
+    ("pos", True, True),
+]
+solve_test_ids = [
+    f'{assume_a}_{"lower" if lower else "upper"}_{"A^T" if transposed else "A"}'
+    for assume_a, lower, transposed in solve_test_cases
+]
+
+
+@pytest.mark.parametrize("b_shape", [(5,), (5, 5)])
+@pytest.mark.parametrize(
+    "assume_a, lower, transposed", solve_test_cases, ids=solve_test_ids
+)
+def test_jax_solve(b_shape: tuple[int], assume_a: str, lower: bool, transposed: bool):
+    A = pt.tensor("A", shape=(5, 5))
+    B = pt.tensor("B", shape=b_shape)
+
+    A_val = np.random.normal(size=(5, 5)).astype(config.floatX)
+    b_val = np.random.normal(size=b_shape).astype(config.floatX)
+
+    if assume_a == "sym":
+        A_val = (A_val + A_val.T) / 2
+    elif assume_a == "pos":
+        A_val = A_val @ A_val.T
+
+    out = pt_slinalg.solve(A, B, assume_a=assume_a, lower=lower, transposed=transposed)
+
+    compare_jax_and_py(
+        graph_inputs=[A, B], graph_outputs=[out], test_inputs=[A_val, b_val]
+    )
+
+
 @pytest.mark.parametrize("check_finite", [False, True])
 @pytest.mark.parametrize("lower", [False, True])
 @pytest.mark.parametrize("trans", [0, 1, 2])