Add graph.duplicate() utility

onnxscript/ir/_core.py

@@ -2307,6 +2307,93 @@ def num_nodes(self) -> int:
         # NOTE: This is a method specific to Graph, not required by the protocol unless proven
         return len(self)
 
+    def duplicate(self) -> Graph:
+        """Create a deep copy of the graph.
+
+        Returns:
+            A new Graph instance that is a deep copy of the current graph.
+            All nodes, values, and relationships are duplicated while maintaining
+            the same structure and connections.
+        """
+        # Create mapping of old values to new values for maintaining connections
+        value_map: dict[Value, Value] = {}
+
+        # Duplicate input values
+        new_inputs = []
+        for input_value in self.inputs:
+            new_input = Value(
+                name=input_value.name,
+                shape=input_value.shape.copy() if input_value.shape else None,
+                type=input_value.type,
+                doc_string=input_value.doc_string,
+                const_value=input_value.const_value,
+            )
+            value_map[input_value] = new_input
+            new_inputs.append(new_input)
+
+        # Duplicate initializers
+        new_initializers = []
+        for init in self.initializers.values():
+            new_init = Value(
+                name=init.name,
+                shape=init.shape.copy() if init.shape else None,
+                type=init.type,
+                doc_string=init.doc_string,
+                const_value=init.const_value,
+            )
+            value_map[init] = new_init
+            new_initializers.append(new_init)
+
+        # Create new graph with inputs and initializers
+        new_graph = Graph(
+            inputs=new_inputs,
+            outputs=[],  # Will be set after nodes are created
+            nodes=[],
+            initializers=new_initializers,
+            doc_string=self.doc_string,
+            opset_imports=dict(self.opset_imports),
+            name=self.name,
+            metadata_props=dict(self.metadata_props) if self.metadata_props else None,
+        )
+
+        # Duplicate nodes while maintaining connections
+        for node in self:
+            # Map old inputs to new inputs using value_map
+            new_inputs = [
+                value_map.get(val) if val is not None else None  # type: ignore[misc]
+                for val in node.inputs
+            ]
+
+            # Create new node
+            new_node = Node(
+                domain=node.domain,
+                op_type=node.op_type,
+                inputs=new_inputs,
+                attributes=list(node.attributes.values()),
+                overload=node.overload,
+                num_outputs=len(node.outputs),
+                version=node.version,
+                graph=new_graph,
+                name=node.name,
+                doc_string=node.doc_string,
+                metadata_props=dict(node.metadata_props) if node.metadata_props else None,
+            )
+
+            # Map old outputs to new outputs
+            for old_output, new_output in zip(node.outputs, new_node.outputs):
+                value_map[old_output] = new_output
+                new_output.name = old_output.name
+                new_output.shape = old_output.shape.copy() if old_output.shape else None
+                new_output.type = old_output.type
+                new_output.doc_string = old_output.doc_string
+                new_output.const_value = old_output.const_value
+
+        # Set graph outputs using mapped values
+        new_outputs = [value_map[output] for output in self.outputs]
+        new_graph.outputs.extend(new_outputs)
+
+        return new_graph
+
     # Mutation methods
     def append(self, node: Node, /) -> None:
         """Append a node to the graph in O(1) time.

onnxscript/ir/_core_test.py

@@ -878,6 +878,43 @@ def setUp(self) -> None:
             opset_imports={"": 1},
         )
 
+    def test_duplicate(self):
+        original_graph = self.graph
+        duplicated_graph = original_graph.duplicate()
+
+        # Test equality of graph properties
+        self.assertEqual(len(original_graph), len(duplicated_graph))
+        self.assertEqual(len(original_graph.inputs), len(duplicated_graph.inputs))
+        self.assertEqual(len(original_graph.outputs), len(duplicated_graph.outputs))
+        self.assertEqual(original_graph.opset_imports, duplicated_graph.opset_imports)
+        self.assertEqual(original_graph.doc_string, duplicated_graph.doc_string)
+        self.assertEqual(original_graph.metadata_props, duplicated_graph.metadata_props)
+        self.assertEqual(original_graph.name, duplicated_graph.name)
+
+        # Verify nodes have same structure but are different instances
+        for orig_node, dup_node in zip(original_graph, duplicated_graph):
+            self.assertEqual(orig_node.op_type, dup_node.op_type)
+            self.assertEqual(orig_node.domain, dup_node.domain)
+            self.assertEqual(orig_node.name, dup_node.name)
+            self.assertEqual(len(orig_node.inputs), len(dup_node.inputs))
+            self.assertEqual(len(orig_node.outputs), len(dup_node.outputs))
+            self.assertNotEqual(id(orig_node), id(dup_node))
+
+        # Verify modifying duplicated graph doesn't affect original
+        new_value = _core.Value(name="add_new")
+        new_node = _core.Node("", "Add", inputs=(new_value,), num_outputs=1)
+        duplicated_graph.append(new_node)
+
+        self.assertEqual(len(original_graph) + 1, len(duplicated_graph))
+        self.assertNotIn(new_node, original_graph)
+
+        # Verify that values are different instances but maintain same properties
+        for orig_val, dup_val in zip(original_graph.inputs, duplicated_graph.inputs):
+            self.assertEqual(orig_val.name, dup_val.name)
+            self.assertEqual(orig_val.shape, dup_val.shape)
+            self.assertEqual(orig_val.type, dup_val.type)
+            self.assertNotEqual(id(orig_val), id(dup_val))
+
     def test_initialize(self):
         self.assertEqual(self.graph.inputs, [self.v0, self.v1])
         self.assertEqual(self.graph.outputs, [*self.node.outputs])