Avoid creating FunctionGraphs for Composite rewrite

Author: ricardoV94

pytensor/scalar/basic.py

@@ -4331,24 +4331,25 @@ def __str__(self):
         if self._name is not None:
             return self._name
 
-        # Rename internal variables
-        for i, r in enumerate(self.fgraph.inputs):
-            r.name = f"i{i}"
-        for i, r in enumerate(self.fgraph.outputs):
-            r.name = f"o{i}"
-        io = set(self.fgraph.inputs + self.fgraph.outputs)
-        for i, r in enumerate(self.fgraph.variables):
-            if (
-                not isinstance(r, Constant)
-                and r not in io
-                and len(self.fgraph.clients[r]) > 1
-            ):
-                r.name = f"t{i}"
+        fgraph = self.fgraph
 
-        if len(self.fgraph.outputs) > 1 or len(self.fgraph.apply_nodes) > 10:
+        if len(fgraph.outputs) > 1 or len(fgraph.apply_nodes) > 10:
             self._name = "Composite{...}"
         else:
-            outputs_str = ", ".join(pprint(output) for output in self.fgraph.outputs)
+            # Rename internal variables
+            for i, r in enumerate(fgraph.inputs):
+                r.name = f"i{i}"
+            for i, r in enumerate(fgraph.outputs):
+                r.name = f"o{i}"
+            io = set(fgraph.inputs + fgraph.outputs)
+            for i, r in enumerate(fgraph.variables):
+                if (
+                    not isinstance(r, Constant)
+                    and r not in io
+                    and len(fgraph.clients[r]) > 1
+                ):
+                    r.name = f"t{i}"
+            outputs_str = ", ".join(pprint(output) for output in fgraph.outputs)
             self._name = f"Composite{{{outputs_str}}}"
 
         return self._name
@@ -4433,9 +4434,10 @@ def c_code_template(self):
         fg = self.fgraph
         subd = {e: f"%(i{i})s" for i, e in enumerate(fg.inputs)}
 
+        inputs_set = frozenset(fg.inputs)
         for var in fg.variables:
             if var.owner is None:
-                if var not in fg.inputs:
+                if var not in inputs_set:
                     # This is an orphan
                     if isinstance(var, Constant) and isinstance(var.type, CLinkerType):
                         subd[var] = f"({var.type.c_literal(var.data)})"

pytensor/tensor/rewriting/elemwise.py

@@ -27,8 +27,9 @@
     out2in,
 )
 from pytensor.graph.rewriting.db import SequenceDB
-from pytensor.graph.traversal import ancestors, toposort
+from pytensor.graph.traversal import ancestors, graph_inputs, toposort
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined
+from pytensor.scalar import ScalarConstant
 from pytensor.scalar.math import Grad2F1Loop, _grad_2f1_loop
 from pytensor.tensor.basic import (
     MakeVector,
@@ -1015,31 +1016,48 @@ def print_profile(stream, prof, level=0):
 @node_rewriter([Elemwise])
 def local_useless_composite_outputs(fgraph, node):
     """Remove inputs and outputs of Composite Ops that are not used anywhere."""
-    if not (
-        isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, ps.Composite)
-    ):
-        return
     comp = node.op.scalar_op
-    used_outputs_idxs = [
-        i for i, o_extern in enumerate(node.outputs) if fgraph.clients[o_extern]
-    ]
-    used_inner_outputs = [comp.outputs[i] for i in used_outputs_idxs]
-    comp_fgraph = FunctionGraph(
-        inputs=comp.inputs, outputs=used_inner_outputs, clone=False
-    )
+
+    if not isinstance(node.op.scalar_op, ps.Composite):
+        return None
+
+    clients = fgraph.clients
+    outer_inputs, outer_outputs = node.inputs, node.outputs
+    inner_inputs, inner_outputs = comp.inputs, comp.outputs
+
+    used_inner_outputs = {
+        inner_out
+        for inner_out, outer_out in zip(inner_outputs, outer_outputs)
+        if clients[outer_out]
+    }
+    used_inner_inputs = {
+        inner_inp
+        for inner_inp in graph_inputs(used_inner_outputs)
+        if not isinstance(inner_inp, ScalarConstant)
+    }
+
+    if len(used_inner_inputs) == len(outer_inputs) or len(used_inner_outputs) == len(
+        outer_outputs
+    ):
+        return None
+
     used_inputs_idxs = [
-        i
-        for i, i_intern in enumerate(comp_fgraph.inputs)
-        if comp_fgraph.clients[i_intern]
+        i for i, inp in enumerate(inner_inputs) if inp in used_inner_inputs
     ]
-    used_inner_inputs = [comp.inputs[i] for i in used_inputs_idxs]
-    if len(used_inner_inputs) < len(node.inputs) or len(used_inner_outputs) < len(
-        node.outputs
-    ):
-        used_inputs = [node.inputs[i] for i in used_inputs_idxs]
-        c = ps.Composite(inputs=used_inner_inputs, outputs=used_inner_outputs)
-        e = Elemwise(scalar_op=c)(*used_inputs, return_list=True)
-        return dict(zip([node.outputs[i] for i in used_outputs_idxs], e, strict=True))
+    used_inner_inputs = [inner_inputs[i] for i in used_inputs_idxs]
+    used_outer_inputs = [outer_inputs[i] for i in used_inputs_idxs]
+
+    new_comp = ps.Composite(inputs=used_inner_inputs, outputs=used_inner_outputs)
+    new_outer_outputs = Elemwise(scalar_op=new_comp)(
+        *used_outer_inputs, return_list=True
+    )
+
+    used_outer_outputs = (
+        outer_outputs[i]
+        for i, out in enumerate(inner_outputs)
+        if out in used_inner_outputs
+    )
+    return dict(zip(used_outer_outputs, new_outer_outputs, strict=True))
 
 
 @node_rewriter([CAReduce])