Fix bug in Composite when multiple outputs are identical

Author: ricardoV94

pytensor/scalar/basic.py

@@ -4146,6 +4146,21 @@ def fgraph(self):
                     "The fgraph to Composite must be exclusively"
                     " composed of ScalarOp instances."
                 )
+
+        # Clone identical outputs that have been merged
+        if len(set(fgraph.outputs)) != len(self.outputs):
+            old_outputs = fgraph.outputs
+            new_outputs = []
+            for output in old_outputs:
+                if output not in new_outputs:
+                    new_outputs.append(output)
+                else:
+                    node = output.owner
+                    output_idx = node.outputs.index(output)
+                    new_output = node.clone().outputs[output_idx]
+                    new_outputs.append(new_output)
+            fgraph = FunctionGraph(fgraph.inputs, new_outputs, clone=False)
+
         self._fgraph = fgraph
         return self._fgraph
 

tests/scalar/test_basic.py

@@ -156,6 +156,17 @@ def test_many_outputs(self):
         fn = make_function(DualLinker().accept(g))
         assert fn(1.0, 2.0, 3.0) == [6.0, 7.0, 0.5]
 
+    def test_identical_outputs(self):
+        x, y, z = floats("xyz")
+        e0 = x + y + z
+        e1 = x + y + z
+        e2 = x / y
+        C = Composite([x, y, z], [e0, e1, e2])
+        c = C.make_node(x, y, z)
+        g = FunctionGraph([x, y, z], c.outputs)
+        fn = make_function(DualLinker().accept(g))
+        assert fn(1.0, 2.0, 3.0) == [6.0, 6.0, 0.5]
+
     def test_composite_printing(self):
         x, y, z = floats("xyz")
         e0 = x + y + z