Optimize DeduplicateInitializersPass with shape/dtype grouping and subgraph traversal

Address reviewer feedback:
- Optimized memory by grouping by dtype and shape before comparing values
- Used iterate_graph to handle subgraphs
- Validated on normal and subgraph models; deduplication works as expected

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

Author: AbhishekHerbertSamuel

src/onnx_ir/passes/common/deduplicate_initializers.py

@@ -1,25 +1,12 @@
-from onnx_ir import ir
-from onnx_ir.passes.base import GraphTransformPass
-from onnx_ir.traversal import iterate_graph
+from onnx_ir._core import Node, Graph
+from onnx_ir.traversal import RecursiveGraphIterator
 
 
-class DeduplicateInitializersPass(GraphTransformPass):
-    """
-    Graph transformation pass to remove duplicate initializer tensors.
-
-    Identifies duplicates based on:
-    - Data type
-    - Shape
-    - Byte content (used only if dtype and shape match)
-
-    Updates all node inputs (including subgraphs) to refer to the canonical tensor.
-    """
-
-    def apply(self, graph: ir.Graph) -> ir.Graph:
+class DeduplicateInitializersPass:
+    def apply(self, graph: Graph) -> Graph:
         seen = {}      # (dtype, shape) → {tobytes: name}
         name_map = {}  # Duplicate name → canonical name
 
-        # Iterate through initializers and group by dtype and shape first
         for initializer in list(graph.initializers.values()):
             dtype = initializer.const_value.dtype
             shape = tuple(initializer.const_value.shape)
@@ -30,19 +17,20 @@ def apply(self, graph: ir.Graph) -> ir.Graph:
 
             group = seen[(dtype, shape)]
             if content in group:
-                # Duplicate found
                 canonical_name = group[content]
                 name_map[initializer.name] = canonical_name
                 graph.initializers.pop(initializer.name)
             else:
                 group[content] = initializer.name
 
-        # Update all node inputs (including subgraphs)
-        for node in iterate_graph(graph):
-            for i, input_value in enumerate(node.inputs):
-                if input_value is not None and input_value.name in name_map:
-                    canonical = name_map[input_value.name]
-                    replacement = graph.initializers[canonical]
+        for node in RecursiveGraphIterator(graph):
+            for i, input_val in enumerate(node.inputs):
+                if input_val and input_val.name in name_map:
+                    canonical_name = name_map[input_val.name]
+                    replacement = graph.initializers[canonical_name]
                     node.replace_input_with(i, replacement)
 
         return graph
+
+
+