Take a graph

Author: justinchuby

onnxscript/ir/_core.py

@@ -1135,7 +1135,7 @@ def __init__(
         num_outputs: int | None = None,
         outputs: Sequence[Value] | None = None,
         version: int | None = None,
-        graph: Graph | None = None,
+        graph: Graph | Function | None = None,
         name: str | None = None,
         doc_string: str | None = None,
         metadata_props: dict[str, str] | None = None,
@@ -1187,7 +1187,7 @@ def __init__(
         self._version: int | None = version
         self._metadata: _metadata.MetadataStore | None = None
         self._metadata_props: dict[str, str] | None = metadata_props
-        self._graph: Graph | None = graph
+        self._graph: Graph | Function | None = graph
         self.doc_string = doc_string
 
         # Add the node as a use of the inputs
@@ -1432,11 +1432,11 @@ def metadata_props(self) -> dict[str, str]:
         return self._metadata_props
 
     @property
-    def graph(self) -> Graph | None:
+    def graph(self) -> Graph | Function | None:
         return self._graph
 
     @graph.setter
-    def graph(self, value: Graph | None) -> None:
+    def graph(self, value: Graph | Function | None) -> None:
         self._graph = value
 
     def op_identifier(self) -> _protocols.OperatorIdentifier:

onnxscript/ir/_tape.py

@@ -44,12 +44,18 @@ class Tape:
             ),
             ir_version=10,
         )
+
+    Attributes:
+        graph_like: The graph to append the new nodes and initializers to. When
+            it is None, the nodes and initializers are creating without owned by a graph.
+            Initializers will not be added to functions because it is not supported by ONNX.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, graph_like: ir.Graph | ir.Function | None = None) -> None:
         self._nodes: list[ir.Node] = []
         self._initializers: list[ir.Value] = []
         self._used_opsets: UsedOpsets = set()
+        self.graph_like = graph_like
 
     def __repr__(self) -> str:
         return f"Tape(nodes={self._nodes}, initializers={self._initializers})"
@@ -92,7 +98,7 @@ def op(
             num_outputs=1,
             overload=overload,
             version=version,
-            graph=graph,
+            graph=graph or self.graph_like,
             name=name,
             doc_string=doc_string,
             metadata_props=metadata_props,
@@ -129,7 +135,7 @@ def op_multi_output(
             num_outputs=num_outputs,
             overload=overload,
             version=version,
-            graph=graph,
+            graph=graph or self.graph_like,
             name=name,
             doc_string=doc_string,
             metadata_props=metadata_props,
@@ -148,6 +154,8 @@ def initializer(self, tensor: ir.TensorProtocol, name: str | None = None) -> ir.
             name=name, shape=shape, type=ir.TensorType(tensor.dtype), const_value=tensor
         )
         self._initializers.append(value)
+        if isinstance(self.graph_like, ir.Graph):
+            self.graph_like.register_initializer(value)
         return value