[IR] Expose the Tape module

onnxscript/ir/__init__.py

@@ -8,6 +8,7 @@
     "traversal",
     "convenience",
     "external_data",
+    "tape",
     # IR classes
     "Tensor",
     "ExternalTensor",
@@ -80,7 +81,7 @@
     "save",
 ]
 
-from onnxscript.ir import convenience, external_data, passes, serde, traversal
+from onnxscript.ir import convenience, external_data, passes, serde, tape, traversal
 from onnxscript.ir._convenience._constructors import node, tensor
 from onnxscript.ir._core import (
     Attr,

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

@@ -2,26 +2,63 @@
 # Licensed under the MIT License.
 """Convenience methods for constructing the IR."""
 
-# NOTE: This is a temporary solution for constructing the IR. It should be replaced
-# with a more permanent solution in the future.
-
 from __future__ import annotations
 
-from typing import Any, Iterable, Iterator, List, Mapping, Optional, Sequence, Tuple
+from typing import (
+    Any,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+)
 
 from onnxscript import ir
 from onnxscript.ir import _convenience
 
+# A type representing the domains/versions used in creating nodes in IR.
+UsedOpsets = set[Tuple[str, Optional[int]]]
+
+
+class Tape:
+    """Tape class.
+
+    A tape is a recorder that collects nodes and initializers that are created so
+    that they can be used for creating a graph.
+
+    Example::
+        from onnxscript import ir
+
+        tape = ir.tape.Tape()
+        a = tape.initializer(ir.tensor([1, 2, 3], name="a"))
+        b: ir.Value = ...
+        c: ir.Value = ...
+        x = tape.op("Add", [a, b], attributes={"alpha": 1.0})
+        y = tape.op("Mul", [x, c], attributes={"beta": 2.0})
+        model = ir.Model(
+            graph := ir.Graph(
+                inputs=[b, c],
+                outputs=[y],
+                nodes=tape.nodes,
+                initializers=tape.initializers
+                opset_imports={"": 20},
+            ),
+            ir_version=10,
+        )
 
-class Tape(Iterable[ir.Node]):
-    """A tape for recording nodes that are created."""
+    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 __iter__(self) -> Iterator[ir.Node]:
-        return iter(self._nodes)
+    def __repr__(self) -> str:
+        return f"Tape(nodes={self._nodes}, initializers={self._initializers})"
 
     @property
     def nodes(self) -> Sequence[ir.Node]:
@@ -31,19 +68,43 @@ def nodes(self) -> Sequence[ir.Node]:
     def initializers(self) -> Sequence[ir.Value]:
         return tuple(self._initializers)
 
+    @property
+    def used_opsets(self) -> UsedOpsets:
+        return self._used_opsets
+
     def op(
         self,
         op_type: str,
         inputs: Sequence[ir.Value | None],
         attributes: Mapping[str, _convenience.SupportedAttrTypes] | None = None,
+        *,
         domain: str = "",
+        overload: str = "",
+        version: int | None = None,
+        graph: ir.Graph | None = None,
+        name: str | None = None,
+        doc_string: str | None = None,
+        metadata_props: dict[str, str] | None = None,
     ) -> ir.Value:
         if attributes is None:
             attrs: Sequence[ir.Attr | ir.RefAttr] = ()
         else:
             attrs = _convenience.convert_attributes(attributes)
-        node = ir.Node(domain, op_type, inputs, attributes=attrs, num_outputs=1)
+        node = ir.Node(
+            domain,
+            op_type,
+            inputs,
+            attributes=attrs,
+            num_outputs=1,
+            overload=overload,
+            version=version,
+            graph=graph or self.graph_like,
+            name=name,
+            doc_string=doc_string,
+            metadata_props=metadata_props,
+        )
         self._nodes.append(node)
+        self._used_opsets.add((domain, version))
 
         return node.outputs[0]
 
@@ -55,13 +116,32 @@ def op_multi_output(
         *,
         num_outputs: int,
         domain: str = "",
+        overload: str = "",
+        version: int | None = None,
+        graph: ir.Graph | None = None,
+        name: str | None = None,
+        doc_string: str | None = None,
+        metadata_props: dict[str, str] | None = None,
     ) -> Sequence[ir.Value]:
         if attributes is None:
             attrs: Sequence[ir.Attr | ir.RefAttr] = ()
         else:
             attrs = _convenience.convert_attributes(attributes)
-        node = ir.Node(domain, op_type, inputs, attributes=attrs, num_outputs=num_outputs)
+        node = ir.Node(
+            domain,
+            op_type,
+            inputs,
+            attributes=attrs,
+            num_outputs=num_outputs,
+            overload=overload,
+            version=version,
+            graph=graph or self.graph_like,
+            name=name,
+            doc_string=doc_string,
+            metadata_props=metadata_props,
+        )
         self._nodes.append(node)
+        self._used_opsets.add((domain, version))
 
         return node.outputs
 
@@ -74,20 +154,14 @@ 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
 
 
-# A type representing the domains/versions used in creating nodes in IR.
-UsedOpsets = List[Tuple[str, Optional[int]]]
-
-
 class Builder(Tape):
     """An extension of the tape that provides a more convenient API for constructing the IR."""
 
-    def __init__(self):
-        super().__init__()
-        self._used_opsets: UsedOpsets = []
-
     def __getattr__(self, op_type: str) -> Any:
         return lambda *args, **kwargs: self._make_node(op_type, args, kwargs)
 
@@ -101,20 +175,22 @@ def _make_node(self, op_type: str, inputs: Sequence[ir.Value], kwargs: dict[str,
             assert isinstance(outputs, int)
             num_outputs = outputs
 
-        self._used_opsets.append((domain, version))
         if num_outputs == 1:
-            value = super().op(op_type, inputs=inputs, attributes=kwargs, domain=domain)
+            value = super().op(
+                op_type, inputs=inputs, attributes=kwargs, domain=domain, version=version
+            )
             if isinstance(outputs, Sequence):
                 value.name = outputs[0]
             return value
         values = super().op_multi_output(
-            op_type, inputs=inputs, attributes=kwargs, domain=domain, num_outputs=num_outputs
+            op_type,
+            inputs=inputs,
+            attributes=kwargs,
+            domain=domain,
+            version=version,
+            num_outputs=num_outputs,
         )
         if isinstance(outputs, Sequence):
             for value, name in zip(values, outputs):
                 value.name = name
         return values
-
-    @property
-    def used_opsets(self) -> UsedOpsets:
-        return self._used_opsets

onnxscript/ir/_tape_test.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import unittest
+
+from onnxscript import ir
+
+
+class TestTape(unittest.TestCase):
+    def test_op(self):
+        # Create a simple ONNX model with shape inference
+        # Define the model
+        inputs = [
+            ir.Value(
+                name="input_a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((1, 2))
+            ),
+            ir.Value(
+                name="input_b", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((1, 2))
+            ),
+        ]
+
+        tape = ir.tape.Tape()
+
+        _ = tape.op("Add", inputs=inputs)
+
+        self.assertEqual([n.op_type for n in tape.nodes], ["Add"])
+
+    def test_initializers(self):
+        inputs = [
+            ir.Value(
+                name="input_a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((1, 2))
+            ),
+            ir.Value(
+                name="input_b",
+                type=ir.TensorType(ir.DataType.FLOAT),
+                shape=ir.Shape((2, 1)),
+                const_value=ir.tensor([[42]] * 2, dtype=ir.DataType.FLOAT),
+            ),
+        ]
+
+        tape = ir.tape.Tape()
+
+        # Shape and type are not explicitly set for the initializer but it should still work
+        initializer = tape.initializer(
+            ir.tensor([[2, 3]], dtype=ir.DataType.FLOAT), name="initializer"
+        )
+        val_add = tape.op("Add", inputs=inputs)
+        _ = tape.op("Mul", inputs=[val_add, initializer])
+
+        self.assertEqual([n.op_type for n in tape.nodes], ["Add", "Mul"])
+        self.assertEqual(tape.initializers, (initializer,))
+
+    def test_op_multi_out(self):
+        inputs = [
+            ir.Value(
+                name="input_a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((1, 2))
+            ),
+            ir.Value(
+                name="input_b",
+                type=ir.TensorType(ir.DataType.FLOAT),
+                shape=ir.Shape((2, 1)),
+                const_value=ir.tensor([[42]] * 2, dtype=ir.DataType.FLOAT),
+            ),
+        ]
+
+        tape = ir.tape.Tape()
+
+        out1, out2, out3 = tape.op_multi_output("SomeOp", inputs=inputs, num_outputs=3)  # pylint: disable=unbalanced-tuple-unpacking
+        _ = tape.op("SomeOtherOp", inputs=[out1, out2, out3])
+
+        self.assertEqual([n.op_type for n in tape.nodes], ["SomeOp", "SomeOtherOp"])
+
+
+if __name__ == "__main__":
+    unittest.main()