onnx-ir 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

onnx_ir/_protocols.py

@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """Protocols for the ONNX IR.
 
 This file defines the interfaces for tools to interact with the IR. The interfaces
@@ -31,18 +31,20 @@ tools.
 from __future__ import annotations
 
 import typing
-from typing import (
-    Any,
+from collections import OrderedDict
+from collections.abc import (
     Collection,
     Iterable,
     Iterator,
     Mapping,
     MutableMapping,
     MutableSequence,
-    OrderedDict,
-    Protocol,
     Sequence,
-    Tuple,
+)
+from typing import (
+    Any,
+    Literal,
+    Protocol,
 )
 
 from onnx_ir import _enums
@@ -52,7 +54,7 @@ if typing.TYPE_CHECKING:
     from typing_extensions import TypeAlias
 
 # An identifier that will uniquely identify an operator. E.g (domain, op_type, overload)
-OperatorIdentifier: TypeAlias = Tuple[str, str, str]
+OperatorIdentifier: TypeAlias = tuple[str, str, str]
 
 
 @typing.runtime_checkable
@@ -277,6 +279,11 @@ class GraphProtocol(Protocol):
     seen as a Sequence of nodes and should be used as such. For example, to obtain
     all nodes as a list, call ``list(graph)``.
 
+    .. :note::
+        ``quantization_annotation`` is deserialized into the Value's ``meta`` field
+        under the ``quant_parameter_tensor_names`` key. Values that are stored
+        under this key will be serialized as quantization annotations.
+
     Attributes:
         name: The name of the graph.
         inputs: The input values of the graph.
@@ -288,7 +295,6 @@ class GraphProtocol(Protocol):
         meta: Metadata store for graph transform passes.
     """
 
-    # TODO(justinchuby): Support quantization_annotation
     name: str | None
     inputs: MutableSequence[ValueProtocol]
     outputs: MutableSequence[ValueProtocol]
@@ -316,11 +322,15 @@ class GraphProtocol(Protocol):
         """Remove a node from the graph."""
         ...
 
-    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_after(
+        self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /
+    ) -> None:
         """Insert new nodes after the given node."""
         ...
 
-    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_before(
+        self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /
+    ) -> None:
         """Insert new nodes before the given node."""
         ...
 
@@ -414,6 +424,8 @@ class AttributeProtocol(Protocol):
     value: Any
     doc_string: str | None
 
+    def is_ref(self) -> Literal[False]: ...
+
 
 @typing.runtime_checkable
 class ReferenceAttributeProtocol(Protocol):
@@ -433,6 +445,8 @@ class ReferenceAttributeProtocol(Protocol):
     type: _enums.AttributeType
     doc_string: str | None
 
+    def is_ref(self) -> Literal[True]: ...
+
 
 @typing.runtime_checkable
 class SparseTensorProtocol(Protocol):
@@ -585,11 +599,15 @@ class FunctionProtocol(Protocol):
         """Remove a node from the function."""
         ...
 
-    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_after(
+        self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /
+    ) -> None:
         """Insert new nodes after the given node."""
         ...
 
-    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_before(
+        self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /
+    ) -> None:
         """Insert new nodes before the given node."""
         ...
 

onnx_ir/_tape.py

@@ -1,77 +1,182 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """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 collections.abc import Mapping, Sequence
+from typing import (
+    Any,
+    Optional,
+)
 
 import onnx_ir as ir
 from onnx_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::
+
+        import onnx_ir as 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]:
         return tuple(self._nodes)
 
+    @property
+    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,
+        output: ir.Value | None = None,
     ) -> ir.Value:
         if attributes is None:
-            attrs: Sequence[ir.Attr | ir.RefAttr] = ()
+            attrs: Sequence[ir.Attr] = ()
         else:
             attrs = _convenience.convert_attributes(attributes)
-        node = ir.Node(domain, op_type, inputs, attributes=attrs, num_outputs=1)
+        output_kwargs: dict[str, Any]
+        if output is None:
+            output_kwargs = dict(num_outputs=1)
+        else:
+            output_kwargs = dict(outputs=[output])
+        node = ir.Node(
+            domain,
+            op_type,
+            inputs,
+            attributes=attrs,
+            **output_kwargs,
+            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]
 
-    def op_multi_output(
+    def op_multi_out(
         self,
         op_type: str,
         inputs: Sequence[ir.Value | None],
         attributes: Mapping[str, _convenience.SupportedAttrTypes] | None = None,
         *,
-        num_outputs: int,
+        num_outputs: int | None = None,
+        outputs: Sequence[ir.Value] | 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,
     ) -> Sequence[ir.Value]:
+        if num_outputs is None and outputs is None:
+            raise ValueError("Either num_outputs or outputs must be provided.")
+        if num_outputs is not None and outputs is not None:
+            raise ValueError("Both num_outputs and outputs cannot be provided simultaneously.")
+        output_kwargs: dict[str, Any]
+        if outputs is None:
+            output_kwargs = dict(num_outputs=num_outputs)
+        else:
+            output_kwargs = dict(outputs=outputs)
         if attributes is None:
-            attrs: Sequence[ir.Attr | ir.RefAttr] = ()
+            attrs: Sequence[ir.Attr] = ()
         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,
+            **output_kwargs,
+            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
 
-
-# A type representing the domains/versions used in creating nodes in IR.
-UsedOpsets = List[Tuple[str, Optional[int]]]
+    def initializer(self, tensor: ir.TensorProtocol, name: str | None = None) -> ir.Value:
+        name = name or tensor.name
+        if name is None:
+            raise ValueError("Name must be provided for initializer.")
+        shape = ir.Shape((d if isinstance(d, int) else d.value) for d in tensor.shape.dims)
+        value = ir.Value(
+            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
 
 
 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)
 
@@ -85,20 +190,22 @@ class Builder(Tape):
             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
+        values = super().op_multi_out(
+            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

onnx_ir/_thirdparty/asciichartpy.py

@@ -1,6 +1,3 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-#
 # Copyright © 2016 Igor Kroitor
 #
 # MIT License
@@ -32,8 +29,8 @@ options to tune the output.
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from math import ceil, floor, isnan
-from typing import Mapping
 
 black = "\033[30m"
 red = "\033[31m"

onnx_ir/_type_casting.py

@@ -1,12 +1,12 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """Numpy utilities for non-native type operation."""
 # TODO(justinchuby): Upstream the logic to onnx
 
 from __future__ import annotations
 
 import typing
-from typing import Sequence
+from collections.abc import Sequence
 
 import ml_dtypes
 import numpy as np
@@ -89,3 +89,18 @@ def unpack_int4(
     """
     unpacked = _unpack_uint4_as_uint8(data, dims)
     return _extend_int4_sign_bits(unpacked).view(ml_dtypes.int4)
+
+
+def unpack_float4e2m1(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[ml_dtypes.float4_e2m1fn]:
+    """Convert a packed float4e2m1 array to unpacked float4e2m1 array.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of float32 reshaped to dims.
+    """
+    return _unpack_uint4_as_uint8(data, dims).view(ml_dtypes.float4_e2m1fn)

onnx_ir/{_internal/version_utils.py → _version_utils.py}

@@ -1,12 +1,9 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """Version utils for testing."""
 
 from __future__ import annotations
 
-import warnings
-from typing import Callable, Sequence
-
 import packaging.version
 
 
@@ -92,27 +89,3 @@ def has_transformers():
         return True  # noqa
     except ImportError:
         return False
-
-
-def ignore_warnings(warns: Warning | Sequence[Warning]) -> Callable:  # type: ignore[arg-type]
-    """Catches warnings.
-
-    Args:
-        warns: warnings to ignore
-
-    Returns:
-        decorated function
-    """
-
-    def wrapper(fct):
-        if warns is None:
-            raise AssertionError(f"warns cannot be None for '{fct}'.")
-
-        def call_f(self):
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore", warns)  # type: ignore[arg-type]
-                return fct(self)
-
-        return call_f
-
-    return wrapper

onnx_ir/convenience.py

@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """Convenience methods for constructing and manipulating the IR."""
 
 from __future__ import annotations
@@ -9,11 +9,13 @@ __all__ = [
     "convert_attributes",
     "replace_all_uses_with",
     "replace_nodes_and_values",
+    "create_value_mapping",
 ]
 
 from onnx_ir._convenience import (
     convert_attribute,
     convert_attributes,
+    create_value_mapping,
     replace_all_uses_with,
     replace_nodes_and_values,
 )