Use actual Solve Op to infer output dtype

CholSolve outputs a different dtype than basic Solve in Scipy==1.15
pymc-devs · Jan 12, 2025 · c7712f6 · c7712f6
1 parent 33c9400
commit c7712f6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/pytensor/tensor/slinalg.py b/pytensor/tensor/slinalg.py
@@ -259,9 +259,10 @@ def make_node(self, A, b):
             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
-        o_dtype = scipy.linalg.solve(
-            np.eye(1).astype(A.dtype), np.eye(1).astype(b.dtype)
-        ).dtype
+        inp_arr = [np.eye(1).astype(A.dtype), np.eye(1).astype(b.dtype)]
+        out_arr = [[None]]
+        self.perform(None, inp_arr, out_arr)
+        o_dtype = out_arr[0][0].dtype
         x = tensor(dtype=o_dtype, shape=b.type.shape)
         return Apply(self, [A, b], [x])
 

diff --git a/tests/tensor/test_slinalg.py b/tests/tensor/test_slinalg.py
@@ -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(
         "A_func, b_func, error_message",
         [
@@ -191,16 +196,16 @@ def test_make_node(self, A_func, b_func, error_message):
         with pytest.raises(ValueError, match=error_message):
             A = A_func()
             b = b_func()
-            SolveBase(b_ndim=2)(A, b)
+            self.SolveTest(b_ndim=2)(A, b)
 
     def test__repr__(self):
         np.random.default_rng(utt.fetch_seed())
         A = matrix()
         b = matrix()
-        y = SolveBase(b_ndim=2)(A, b)
+        y = self.SolveTest(b_ndim=2)(A, b)
         assert (
             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"
         )
 
 
@@ -450,7 +455,7 @@ def test_solve_dtype(self):
             fn = function([A, b], x)
             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():