Fix bug in local_dimshuffle_lift when elemwise has multiple outputs

Author: ricardoV94

pytensor/tensor/rewriting/elemwise.py

@@ -422,7 +422,12 @@ def local_dimshuffle_lift(fgraph, node):
     inp = node.inputs[0]
     inode = inp.owner
     new_order = op.new_order
-    if inode and isinstance(inode.op, Elemwise) and (len(fgraph.clients[inp]) == 1):
+    if (
+        inode
+        and isinstance(inode.op, Elemwise)
+        and len(inode.outputs) == 1
+        and (len(fgraph.clients[inp]) == 1)
+    ):
         # Don't use make_node to have tag.test_value set.
         new_inputs = []
         for inp in inode.inputs:

tests/tensor/rewriting/test_elemwise.py

@@ -19,7 +19,7 @@
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.misc.safe_asarray import _asarray
 from pytensor.raise_op import assert_op
-from pytensor.scalar.basic import Composite
+from pytensor.scalar.basic import Composite, float64
 from pytensor.tensor.basic import MakeVector
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.math import abs as at_abs
@@ -163,6 +163,20 @@ def test_dimshuffle_on_broadcastable(self):
         # Check stacktrace was copied over correctly after rewrite was applied
         assert hasattr(g.outputs[0].tag, "trace")
 
+    def test_dimshuffle_lift_multi_out_elemwise(self):
+        # Create a multi-output Elemwise Op with Composite
+        x = float64("x")
+        outs = [x + 1, x + 2]
+        op = Elemwise(Composite([x], outs))
+
+        # Transpose both outputs
+        x = matrix("x")
+        outs = [out.T for out in op(x)]
+
+        # Make sure rewrite doesn't apply in this case
+        g = FunctionGraph([x], outs)
+        assert not local_dimshuffle_lift.transform(g, g.outputs[0].owner)
+
 
 def test_local_useless_dimshuffle_in_reshape():
     vec = TensorType(dtype="float64", shape=(None,))("vector")