Addressed feedback by ricardoV94

Author: Luca Citi

pytensor/graph/rewriting/basic.py

@@ -1636,6 +1636,10 @@ def transform(self, fgraph, node, get_nodes=True):
         if node.op != self.op:
             return False
 
+        if len(node.outputs) != 1:
+            # PatternNodeRewriter doesn't support replacing multi-output nodes
+            return False
+
         s = unify(self.in_pattern, node.out)
 
         if s is False:
@@ -1658,27 +1662,20 @@ def transform(self, fgraph, node, get_nodes=True):
             ):
                 return False
 
-        if ret.owner:
-            if len(node.outputs) != len(ret.owner.outputs):
-                return
-            if len(node.outputs) > 1:
-                if not all(
-                    o.type.is_super(new_o.type)
-                    for o, new_o in zip(node.outputs, ret.owner.outputs, strict=True)
-                ):
-                    return False
-            else:
-                if self.allow_cast:
-                    out_dtype = getattr(node.outputs[0].type, "dtype", None)
-                    ret_dtype = getattr(ret.owner.outputs[0].type, "dtype", None)
-                    if ret_dtype != out_dtype:
-                        ret = pytensor.tensor.basic.cast(ret, out_dtype)
-                if not node.outputs[0].type.is_super(ret.owner.outputs[0].type):
-                    return False
-        else:
-            # ret is just an input variable
-            assert len(node.outputs) == 1
-            if not node.outputs[0].type.is_super(ret.type):
+        [old_out] = node.outputs
+        if not old_out.type.is_super(ret.type):
+            # Type doesn't match
+            if not (
+                self.allow_cast
+                and isinstance(old_out.type, pytensor.tensor.TensorType)
+                and isinstance(ret.type, pytensor.tensor.TensorType)
+            ):
+                return False
+
+            # Try to cast tensors
+            ret = ret.astype(old_out.type.dtype)
+            if not old_out.type.is_super(ret.type):
+                # Still doesn't match
                 return False
 
         return [ret]

tests/tensor/rewriting/test_math.py

@@ -4142,16 +4142,10 @@ def test_local_1msigmoid(self):
             (np.float32(1 - 9e-6) - sigmoid(x), np.float32(1 - 9e-6) - sigmoid(x)),
             (np.float64(1 - 1e-9) - sigmoid(xd), np.float64(1 - 1e-9) - sigmoid(xd)),
         ]:
-            f = pytensor.function([x, xd], out, m, on_unused_input="ignore")
-            f_outs = f.maker.fgraph.outputs
-            assert equal_computations(
-                f_outs, [expected]
-            ), "Expression:\n{}rewritten as:\n{}expected:\n{}".format(
-                *(
-                    pytensor.dprint(expr, print_type=True, file="str")
-                    for expr in (out, f_outs, expected)
-                )
+            rewritten = rewrite_graph(
+                out, include=["canonicalize", "specialize", "stabilize"]
             )
+            utt.assert_equal_computations([rewritten], [expected], original=out)
 
     def test_local_sigm_times_exp(self):
         """

tests/unittest_tools.py

@@ -11,6 +11,7 @@
 from pytensor.compile.debugmode import str_diagnostic
 from pytensor.configdefaults import config
 from pytensor.gradient import verify_grad as orig_verify_grad
+from pytensor.graph.basic import equal_computations
 from pytensor.tensor.basic import as_tensor_variable
 from pytensor.tensor.math import _allclose
 from pytensor.tensor.math import add as pt_add
@@ -279,6 +280,41 @@ def assert_allclose(expected, value, rtol=None, atol=None):
         raise WrongValue(expected, value, rtol, atol)
 
 
+def assert_equal_computations(rewritten, expected, *args, original=None, **kwargs):
+    """
+    Assert that `rewritten` computes the same as `expected`.
+
+    Parameters
+    ----------
+    rewritten
+        The expression after the rewrite pass.
+    expected
+        The reference expression to compare against.
+    *args, **kwargs
+        Extra arguments forwarded to equal_computations.
+    original : optional
+        If given, will be printed in the error message.
+    """
+    __tracebackhide__ = True  # Hide traceback for py.test
+
+    ok = equal_computations(rewritten, expected, *args, **kwargs)
+
+    if not ok:
+        parts = []
+
+        def _dprint(expr):
+            return pytensor.dprint(expr, print_type=True, file="str")
+
+        if original is not None:
+            parts.append(f"\nOriginal:\n{_dprint(original)}")
+        parts.append(f"\nRewritten:\n{_dprint(rewritten)}")
+        parts.append(f"\nExpected:\n{_dprint(expected)}")
+
+        raise AssertionError("equal_computations failed\n" + "".join(parts))
+
+    return True
+
+
 class AttemptManyTimes:
     """Decorator for unit tests that forces a unit test to be attempted
     multiple times. The test needs to pass a certain number of times for it to