fix: make elemwise test check against dtype in the graph

pymc-devs · Aug 13, 2024 · fe0a7ec · fe0a7ec
1 parent 584e506
commit fe0a7ec
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/tensor/utils.py b/tests/tensor/utils.py
@@ -508,15 +508,17 @@ def test_good(self):
                 if not isinstance(expecteds, list | tuple):
                     expecteds = (expecteds,)
 
-                for i, (variable, expected) in enumerate(zip(variables, expecteds)):
+                for i, (variable, expected, out_symbol) in enumerate(
+                    zip(variables, expecteds, node.outputs)
+                ):
                     condition = (
-                        variable.dtype != expected.dtype
+                        variable.dtype != out_symbol.type.dtype
                         or variable.shape != expected.shape
                         or not np.allclose(variable, expected, atol=eps, rtol=eps)
                     )
                     assert not condition, (
                         f"Test {self.op}::{testname}: Output {i} gave the wrong"
-                        f" value. With inputs {inputs}, expected {expected} (dtype {expected.dtype}),"
+                        f" value. With inputs {inputs}, expected {expected} (dtype {out_symbol.type.dtype}),"
                         f" got {variable} (dtype {variable.dtype}). eps={eps:f}"
                         f" np.allclose returns {np.allclose(variable, expected, atol=eps, rtol=eps)} {np.allclose(variable, expected)}"
                     )