Remove jax tests

Author: jessegrabowski

tests/link/jax/test_slinalg.py

@@ -104,33 +104,16 @@ def test_jax_basic():
     )
 
 
-@pytest.mark.parametrize(
-    "b_shape",
-    [(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("lower", [False, True])
-@pytest.mark.parametrize("transposed", [False, True])
-def test_jax_solve(b_shape: tuple[int], assume_a, lower, transposed):
+def test_jax_solve():
     rng = np.random.default_rng(utt.fetch_seed())
 
     A = pt.tensor("A", shape=(5, 5))
-    b = pt.tensor("B", shape=b_shape)
-
-    def A_func(x):
-        if assume_a == "sym":
-            return (x + x.T) / 2
-        if assume_a == "pos":
-            return x @ x.T
-        return x
+    b = pt.tensor("B", shape=(5, 5))
 
-    out = pt_slinalg.solve(
-        A_func(A), b, assume_a=assume_a, lower=lower, transposed=transposed
-    )
+    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=b_shape).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
 
     compare_jax_and_py(
         [A, b],
@@ -139,35 +122,21 @@ def A_func(x):
     )
 
 
-@pytest.mark.parametrize(
-    "b_shape", [(5, 1), (5, 5), (5,)], ids=["b_col_vec", "b_matrix", "b_vec"]
-)
-@pytest.mark.parametrize("lower", [False, True])
-@pytest.mark.parametrize("trans", [0, 1, 2])
-@pytest.mark.parametrize("unit_diagonal", [False, True])
-def test_jax_SolveTriangular(b_shape: tuple[int], lower, trans, unit_diagonal):
+def test_jax_SolveTriangular():
     rng = np.random.default_rng(utt.fetch_seed())
 
     A = pt.tensor("A", shape=(5, 5))
-    b = pt.tensor("B", shape=b_shape)
-
-    def A_func(x):
-        x = x @ x.T
-        x = pt.linalg.cholesky(x, lower=lower)
-        if unit_diagonal:
-            x = pt.fill_diagonal(x, 1.0)
-
-        return x
+    b = pt.tensor("B", shape=(5, 5))
 
     A_val = rng.normal(size=(5, 5)).astype(config.floatX)
-    b_val = rng.normal(size=b_shape).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
 
     out = pt_slinalg.solve_triangular(
-        A_func(A),
+        A,
         b,
-        trans=trans,
-        lower=lower,
-        unit_diagonal=unit_diagonal,
+        trans=0,
+        lower=True,
+        unit_diagonal=False,
     )
     compare_jax_and_py([A, b], [out], [A_val, b_val])