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,72 @@ def test_jax_basic():
  )
 
 
-@pytest.mark.parametrize("check_finite", [False, True])
+@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):
+ 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
+
+ out = pt_slinalg.solve(
+ A_func(A), b, assume_a=assume_a, lower=lower, transposed=transposed
+ )
+
+ A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+ b_val = rng.normal(size=b_shape).astype(config.floatX)
+
+ compare_jax_and_py(
+ [A, b],
+ [out],
+ [A_val, b_val],
+ )
+
+
+@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])
-def test_jax_SolveTriangular(trans, lower, check_finite):
- x = matrix("x")
- b = vector("b")
+@pytest.mark.parametrize("unit_diagonal", [False, True])
+def test_jax_SolveTriangular(b_shape: tuple[int], lower, trans, unit_diagonal):
+ 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
+
+ A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+ b_val = rng.normal(size=b_shape).astype(config.floatX)
 
  out = pt_slinalg.solve_triangular(
- x,
+ A_func(A),
  b,
  trans=trans,
  lower=lower,
- check_finite=check_finite,
- )
- compare_jax_and_py(
- [x, b],
- [out],
- [
- np.eye(10).astype(config.floatX),
- np.arange(10).astype(config.floatX),
- ],
+ unit_diagonal=unit_diagonal,
  )
+ 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],
  )