提交 581f65a1 authored 作者: ricardoV94's avatar ricardoV94 提交者: Ricardo Vieira

Adapt to Solve changes in Scipy 1.15

1. Use actual Solve Op to infer output dtype as CholSolve outputs a different dtype than basic Solve in Scipy==1.15 2. Tweaked test related to https://github.com/pymc-devs/pytensor/issues/1152 3. Tweak tolerage
上级 cff058c9
...@@ -259,9 +259,10 @@ class SolveBase(Op): ...@@ -259,9 +259,10 @@ class SolveBase(Op):
raise ValueError(f"`b` must have {self.b_ndim} dims; got {b.type} instead.") raise ValueError(f"`b` must have {self.b_ndim} dims; got {b.type} instead.")
# Infer dtype by solving the most simple case with 1x1 matrices # Infer dtype by solving the most simple case with 1x1 matrices
o_dtype = scipy.linalg.solve( inp_arr = [np.eye(1).astype(A.dtype), np.eye(1).astype(b.dtype)]
np.eye(1).astype(A.dtype), np.eye(1).astype(b.dtype) out_arr = [[None]]
).dtype self.perform(None, inp_arr, out_arr)
o_dtype = out_arr[0][0].dtype
x = tensor(dtype=o_dtype, shape=b.type.shape) x = tensor(dtype=o_dtype, shape=b.type.shape)
return Apply(self, [A, b], [x]) return Apply(self, [A, b], [x])
......
...@@ -590,7 +590,7 @@ class TestInplace: ...@@ -590,7 +590,7 @@ class TestInplace:
A_val_copy, b_val_copy A_val_copy, b_val_copy
) )
np.testing.assert_allclose( np.testing.assert_allclose(
out, expected_out, atol=1e-5 if config.floatX == "float32" else 0 out, expected_out, atol=1e-4 if config.floatX == "float32" else 0
) )
# Confirm input was destroyed # Confirm input was destroyed
......
...@@ -169,7 +169,12 @@ def test_eigvalsh_grad(): ...@@ -169,7 +169,12 @@ def test_eigvalsh_grad():
) )
class TestSolveBase(utt.InferShapeTester): class TestSolveBase:
class SolveTest(SolveBase):
def perform(self, node, inputs, outputs):
A, b = inputs
outputs[0][0] = scipy.linalg.solve(A, b)
@pytest.mark.parametrize( @pytest.mark.parametrize(
"A_func, b_func, error_message", "A_func, b_func, error_message",
[ [
...@@ -191,16 +196,16 @@ class TestSolveBase(utt.InferShapeTester): ...@@ -191,16 +196,16 @@ class TestSolveBase(utt.InferShapeTester):
with pytest.raises(ValueError, match=error_message): with pytest.raises(ValueError, match=error_message):
A = A_func() A = A_func()
b = b_func() b = b_func()
SolveBase(b_ndim=2)(A, b) self.SolveTest(b_ndim=2)(A, b)
def test__repr__(self): def test__repr__(self):
np.random.default_rng(utt.fetch_seed()) np.random.default_rng(utt.fetch_seed())
A = matrix() A = matrix()
b = matrix() b = matrix()
y = SolveBase(b_ndim=2)(A, b) y = self.SolveTest(b_ndim=2)(A, b)
assert ( assert (
y.__repr__() y.__repr__()
== "SolveBase{lower=False, check_finite=True, b_ndim=2, overwrite_a=False, overwrite_b=False}.0" == "SolveTest{lower=False, check_finite=True, b_ndim=2, overwrite_a=False, overwrite_b=False}.0"
) )
...@@ -239,8 +244,9 @@ class TestSolve(utt.InferShapeTester): ...@@ -239,8 +244,9 @@ class TestSolve(utt.InferShapeTester):
A_val = np.asarray(rng.random((5, 5)), dtype=config.floatX) A_val = np.asarray(rng.random((5, 5)), dtype=config.floatX)
A_val = np.dot(A_val.transpose(), A_val) A_val = np.dot(A_val.transpose(), A_val)
assert np.allclose( np.testing.assert_allclose(
scipy.linalg.solve(A_val, b_val), gen_solve_func(A_val, b_val) scipy.linalg.solve(A_val, b_val, assume_a="gen"),
gen_solve_func(A_val, b_val),
) )
A_undef = np.array( A_undef = np.array(
...@@ -253,7 +259,7 @@ class TestSolve(utt.InferShapeTester): ...@@ -253,7 +259,7 @@ class TestSolve(utt.InferShapeTester):
], ],
dtype=config.floatX, dtype=config.floatX,
) )
assert np.allclose( np.testing.assert_allclose(
scipy.linalg.solve(A_undef, b_val), gen_solve_func(A_undef, b_val) scipy.linalg.solve(A_undef, b_val), gen_solve_func(A_undef, b_val)
) )
...@@ -450,7 +456,7 @@ class TestCholeskySolve(utt.InferShapeTester): ...@@ -450,7 +456,7 @@ class TestCholeskySolve(utt.InferShapeTester):
fn = function([A, b], x) fn = function([A, b], x)
x_result = fn(A_val.astype(A_dtype), b_val.astype(b_dtype)) x_result = fn(A_val.astype(A_dtype), b_val.astype(b_dtype))
assert x.dtype == x_result.dtype assert x.dtype == x_result.dtype, (A_dtype, b_dtype)
def test_cho_solve(): def test_cho_solve():
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论