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

onnx_ir/__init__.py

@@ -1,14 +1,19 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """In-memory intermediate representation for ONNX graphs."""
 
 __all__ = [
     # Modules
     "serde",
+    "traversal",
+    "convenience",
+    "external_data",
+    "tape",
     # IR classes
     "Tensor",
     "ExternalTensor",
     "StringTensor",
+    "LazyTensor",
     "SymbolicDim",
     "Shape",
     "TensorType",
@@ -66,19 +71,24 @@ __all__ = [
     "TensorProtoTensor",
     # Conversion functions
     "from_proto",
+    "from_onnx_text",
     "to_proto",
-    # IR Tensor initializer
+    # Convenience constructors
     "tensor",
+    "node",
     # Pass infrastructure
     "passes",
-    "traversal",
     # IO
     "load",
     "save",
+    # Flags
+    "DEBUG",
 ]
 
-from onnx_ir import passes, serde, traversal
-from onnx_ir._convenience import tensor
+import types
+
+from onnx_ir import convenience, external_data, passes, serde, tape, traversal
+from onnx_ir._convenience._constructors import node, tensor
 from onnx_ir._core import (
     Attr,
     AttrFloat32,
@@ -100,6 +110,7 @@ from onnx_ir._core import (
     Graph,
     GraphView,
     Input,
+    LazyTensor,
     Model,
     Node,
     OptionalType,
@@ -138,17 +149,19 @@ from onnx_ir._protocols import (
     TypeProtocol,
     ValueProtocol,
 )
-from onnx_ir.serde import TensorProtoTensor, from_proto, to_proto
-
+from onnx_ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_proto
 
-DEBUG: bool = False
+DEBUG = False
 
 
 def __set_module() -> None:
     """Set the module of all functions in this module to this public module."""
     global_dict = globals()
     for name in __all__:
-        global_dict[name].__module__ = __name__
+        obj = global_dict[name]
+        if hasattr(obj, "__module__") and not isinstance(obj, types.GenericAlias):
+            obj.__module__ = __name__
 
 
 __set_module()
+__version__ = "0.1.0"

onnx_ir/{_convenience.py → _convenience/__init__.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.
 
 This is an internal only module. We should choose to expose some of the methods
@@ -12,19 +12,17 @@ __all__ = [
     "convert_attribute",
     "convert_attributes",
     "replace_all_uses_with",
+    "create_value_mapping",
+    "replace_nodes_and_values",
 ]
 
-import typing
-from typing import Mapping, Sequence, Union
+from collections.abc import Mapping, Sequence
+from typing import Union
 
-import numpy as np
 import onnx
 
 from onnx_ir import _core, _enums, _protocols, serde
 
-if typing.TYPE_CHECKING:
-    import numpy.typing as npt
-
 SupportedAttrTypes = Union[
     str,
     int,
@@ -35,9 +33,9 @@ SupportedAttrTypes = Union[
     _protocols.TensorProtocol,  # This includes all in-memory tensor types
     onnx.TensorProto,
     _core.Attr,
-    _core.RefAttr,
     _protocols.GraphProtocol,
     Sequence[_protocols.GraphProtocol],
+    onnx.GraphProto,
     _protocols.TypeProtocol,
     Sequence[_protocols.TypeProtocol],
     None,
@@ -52,7 +50,7 @@ def _infer_attribute_type(attr: SupportedAttrTypes) -> _enums.AttributeType:
         return _enums.AttributeType.FLOAT
     if isinstance(attr, str):
         return _enums.AttributeType.STRING
-    if isinstance(attr, (_core.Attr, _core.RefAttr)):
+    if isinstance(attr, _core.Attr):
         return attr.type
     if isinstance(attr, Sequence) and all(isinstance(x, int) for x in attr):
         return _enums.AttributeType.INTS
@@ -63,10 +61,15 @@ def _infer_attribute_type(attr: SupportedAttrTypes) -> _enums.AttributeType:
     if isinstance(attr, (_core.TensorBase, onnx.TensorProto, _protocols.TensorProtocol)):
         # Be sure to check TensorProtocol last because isinstance checking on Protocols can be slower
         return _enums.AttributeType.TENSOR
-    if isinstance(attr, (_core.Graph, _protocols.GraphProtocol)):
+    if isinstance(attr, Sequence) and all(
+        isinstance(x, (_core.TensorBase, onnx.TensorProto, _protocols.TensorProtocol))
+        for x in attr
+    ):
+        return _enums.AttributeType.TENSORS
+    if isinstance(attr, (_core.Graph, onnx.GraphProto, _protocols.GraphProtocol)):
         return _enums.AttributeType.GRAPH
     if isinstance(attr, Sequence) and all(
-        isinstance(x, (_core.Graph, _protocols.GraphProtocol)) for x in attr
+        isinstance(x, (_core.Graph, onnx.GraphProto, _protocols.GraphProtocol)) for x in attr
     ):
         return _enums.AttributeType.GRAPHS
     if isinstance(
@@ -94,7 +97,7 @@ def convert_attribute(
     name: str,
     attr: SupportedAttrTypes,
     attr_type: _enums.AttributeType | None = None,
-) -> _core.Attr | _core.RefAttr:
+) -> _core.Attr:
     """Convert a Python object to a _core.Attr object.
 
     This method is useful when constructing nodes with attributes. It infers the
@@ -110,7 +113,7 @@ def convert_attribute(
         A ``Attr`` object.
 
     Raises:
-        ValueError: If :param:`attr` is ``None`` and :param:`attr_type` is not provided.
+        ValueError: If ``attr`` is ``None`` and ``attr_type`` is not provided.
         TypeError: If the type of the attribute is not supported.
     """
     if attr is None:
@@ -118,7 +121,7 @@ def convert_attribute(
             raise ValueError("attr_type must be provided when attr is None")
         return _core.Attr(name, attr_type, None)
 
-    if isinstance(attr, (_core.Attr, _core.RefAttr)):
+    if isinstance(attr, _core.Attr):
         if attr.name != name:
             raise ValueError(
                 f"Attribute name '{attr.name}' does not match provided name '{name}'"
@@ -148,11 +151,27 @@ def convert_attribute(
         if isinstance(attr, (_core.TensorBase, _protocols.TensorProtocol)):
             return _core.AttrTensor(name, attr)
         if isinstance(attr, onnx.TensorProto):
-            return _core.AttrTensor(name, serde.TensorProtoTensor(attr))
+            return _core.AttrTensor(name, serde.deserialize_tensor(attr))
+    if attr_type == _enums.AttributeType.TENSORS:
+        tensors = []
+        for t in attr:  # type: ignore[union-attr]
+            if isinstance(t, onnx.TensorProto):
+                tensors.append(_core.AttrTensor(name, serde.deserialize_tensor(t)))
+            else:
+                tensors.append(t)  # type: ignore[arg-type]
+        return _core.AttrTensors(name, tensors)  # type: ignore[arg-type]
     if attr_type == _enums.AttributeType.GRAPH:
+        if isinstance(attr, onnx.GraphProto):
+            attr = serde.deserialize_graph(attr)
         return _core.AttrGraph(name, attr)  # type: ignore[arg-type]
     if attr_type == _enums.AttributeType.GRAPHS:
-        return _core.AttrGraphs(name, attr)  # type: ignore[arg-type]
+        graphs = []
+        for graph in attr:  # type: ignore[union-attr]
+            if isinstance(graph, onnx.GraphProto):
+                graphs.append(serde.deserialize_graph(graph))
+            else:
+                graphs.append(graph)  # type: ignore[arg-type]
+        return _core.AttrGraphs(name, graphs)  # type: ignore[arg-type]
     if attr_type == _enums.AttributeType.TYPE_PROTO:
         return _core.AttrTypeProto(name, attr)  # type: ignore[arg-type]
     if attr_type == _enums.AttributeType.TYPE_PROTOS:
@@ -162,7 +181,7 @@ def convert_attribute(
 
 def convert_attributes(
     attrs: Mapping[str, SupportedAttrTypes],
-) -> list[_core.Attr | _core.RefAttr]:
+) -> list[_core.Attr]:
     """Convert a dictionary of attributes to a list of _core.Attr objects.
 
     It infers the attribute type based on the type of the value. The supported
@@ -193,7 +212,7 @@ def convert_attributes(
         ...     "type_protos": [ir.TensorType(ir.DataType.FLOAT), ir.TensorType(ir.DataType.FLOAT)],
         ... }
         >>> convert_attributes(attrs)
-        [Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(name='proto')), Attr('graph', INTS, Graph(
+        [Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', INTS, Graph(
             name='graph0',
             inputs=(
         <BLANKLINE>
@@ -228,7 +247,7 @@ def convert_attributes(
     Returns:
         A list of _core.Attr objects.
     """
-    attributes: list[_core.Attr | _core.RefAttr] = []
+    attributes: list[_core.Attr] = []
     for name, attr in attrs.items():
         if attr is not None:
             attributes.append(convert_attribute(name, attr))
@@ -291,86 +310,6 @@ def replace_all_uses_with(
             user_node.replace_input_with(index, replacement)
 
 
-def tensor(
-    value: npt.ArrayLike
-    | onnx.TensorProto
-    | _protocols.DLPackCompatible
-    | _protocols.ArrayCompatible,
-    dtype: _enums.DataType | None = None,
-    name: str | None = None,
-    doc_string: str | None = None,
-) -> _protocols.TensorProtocol:
-    """Create a tensor value from an ArrayLike object or a TensorProto.
-
-    The dtype must match the value. Reinterpretation of the value is
-    not supported, unless if the value is a plain Python object, in which case
-    it is converted to a numpy array with the given dtype.
-
-    :param:`value` can be a numpy array, a plain Python object, or a TensorProto.
-
-    Example::
-
-        >>> import onnx_ir as ir
-        >>> import numpy as np
-        >>> import ml_dtypes
-        >>> import onnx
-        >>> ir.tensor(np.array([1, 2, 3], dtype=np.int16))
-        Tensor<INT16,[3]>(array([1, 2, 3], dtype=int16), name=None)
-        >>> ir.tensor([1, 2, 3], dtype=ir.DataType.BFLOAT16)
-        Tensor<BFLOAT16,[3]>(array([1, 2, 3], dtype=bfloat16), name=None)
-        >>> tp_tensor = ir.tensor(onnx.helper.make_tensor("tensor", onnx.TensorProto.FLOAT, dims=[], vals=[0.5]))
-        >>> tp_tensor.numpy()
-        array(0.5, dtype=float32)
-
-    Args:
-        value: The numpy array to create the tensor from.
-        dtype: The data type of the tensor.
-        name: The name of the tensor.
-        doc_string: The documentation string of the tensor.
-
-    Returns:
-        A tensor value.
-
-    Raises:
-        ValueError: If the dtype does not match the value when value is not a plain Python
-            object like ``list[int]``.
-    """
-    if isinstance(value, _protocols.TensorProtocol):
-        if dtype is not None and dtype != value.dtype:
-            raise ValueError(
-                f"The dtype must match the value when value is a Tensor. dtype={dtype}, value.dtype={value.dtype}. "
-                "You do not have to specify the dtype when value is a Tensor."
-            )
-        return value
-    if isinstance(value, onnx.TensorProto):
-        tensor_ = serde.deserialize_tensor(value)
-        if name is not None:
-            tensor_.name = name
-        if doc_string is not None:
-            tensor_.doc_string = doc_string
-        if dtype is not None and dtype != tensor_.dtype:
-            raise ValueError(
-                f"The dtype must match the value when value is a TensorProto. dtype={dtype}, value.data_type={tensor_.dtype}"
-                "You do not have to specify the dtype when value is a TensorProto."
-            )
-    elif isinstance(value, (_protocols.DLPackCompatible, _protocols.ArrayCompatible)):
-        tensor_ = _core.Tensor(value, dtype=dtype, name=name, doc_string=name)
-    else:
-        if dtype is not None:
-            numpy_dtype = dtype.numpy()
-        else:
-            numpy_dtype = None
-        array = np.array(value, dtype=numpy_dtype)
-        tensor_ = _core.Tensor(
-            array,
-            dtype=dtype,
-            shape=_core.Shape(array.shape),
-            name=name,
-            doc_string=name,
-        )
-    return tensor_
-
-
 def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     """Return a dictionary mapping names to values in the graph.
 
@@ -382,7 +321,7 @@ def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     Returns:
         A dictionary mapping names to values.
     """
-    values = {}
+    values: dict[str, _core.Value] = {}
     values.update(graph.initializers)
     # The names of the values can be None or "", which we need to exclude
     for input in graph.inputs:
@@ -416,7 +355,6 @@ def replace_nodes_and_values(
         old_values: The values to replace.
         new_values: The values to replace with.
     """
-
     for old_value, new_value in zip(old_values, new_values):
         # Propagate relevant info from old value to new value
         # TODO(Rama): Perhaps this should be a separate utility function. Also, consider

onnx_ir/_convenience/_constructors.py

@@ -0,0 +1,213 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Convenience constructors for IR objects."""
+
+from __future__ import annotations
+
+__all__ = [
+    "tensor",
+    "node",
+]
+
+import typing
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+import onnx
+
+from onnx_ir import _convenience, _core, _enums, _protocols, serde, tensor_adapters
+
+if typing.TYPE_CHECKING:
+    import numpy.typing as npt
+
+    import onnx_ir as ir
+
+
+def tensor(
+    value: npt.ArrayLike | onnx.TensorProto | ir.DLPackCompatible | ir.ArrayCompatible,
+    dtype: _enums.DataType | None = None,
+    name: str | None = None,
+    doc_string: str | None = None,
+) -> _protocols.TensorProtocol:
+    """Create a tensor value from an ArrayLike object or a TensorProto.
+
+    The dtype must match the value. Reinterpretation of the value is
+    not supported, unless if the value is a plain Python object, in which case
+    it is converted to a numpy array with the given dtype.
+
+    ``value`` can be a numpy array, a plain Python object, or a TensorProto.
+
+    Example::
+
+        >>> import onnx_ir as ir
+        >>> import numpy as np
+        >>> import ml_dtypes
+        >>> import onnx
+        >>> ir.tensor(np.array([1, 2, 3], dtype=np.int16))
+        Tensor<INT16,[3]>(array([1, 2, 3], dtype=int16), name=None)
+        >>> ir.tensor([1, 2, 3], dtype=ir.DataType.BFLOAT16)
+        Tensor<BFLOAT16,[3]>(array([1, 2, 3], dtype=bfloat16), name=None)
+        >>> tp_tensor = ir.tensor(onnx.helper.make_tensor("tensor", onnx.TensorProto.FLOAT, dims=[], vals=[0.5]))
+        >>> tp_tensor.numpy()
+        array(0.5, dtype=float32)
+        >>> import torch
+        >>> ir.tensor(torch.tensor([1.0, 2.0]), name="torch_tensor")
+        TorchTensor<FLOAT,[2]>(tensor([1., 2.]), name='torch_tensor')
+
+    Args:
+        value: The numpy array to create the tensor from.
+        dtype: The data type of the tensor.
+        name: The name of the tensor.
+        doc_string: The documentation string of the tensor.
+
+    Returns:
+        A tensor value.
+
+    Raises:
+        ValueError: If the dtype does not match the value when value is not a plain Python
+            object like ``list[int]``.
+    """
+    if isinstance(value, _protocols.TensorProtocol):
+        if dtype is not None and dtype != value.dtype:
+            raise ValueError(
+                f"The dtype must match the value when value is a Tensor. dtype={dtype}, value.dtype={value.dtype}. "
+                "You do not have to specify the dtype when value is a Tensor."
+            )
+        return value
+    if isinstance(value, onnx.TensorProto):
+        tensor_ = serde.deserialize_tensor(value)
+        if name is not None:
+            tensor_.name = name
+        if doc_string is not None:
+            tensor_.doc_string = doc_string
+        if dtype is not None and dtype != tensor_.dtype:
+            raise ValueError(
+                f"The dtype must match the value when value is a TensorProto. dtype={dtype}, value.data_type={tensor_.dtype}"
+                "You do not have to specify the dtype when value is a TensorProto."
+            )
+        return tensor_
+    elif str(type(value)) == "<class 'torch.Tensor'>":
+        # NOTE: We use str(type(...)) and do not import torch for type checking
+        # as it creates overhead during import
+        return tensor_adapters.TorchTensor(value, name=name, doc_string=doc_string)  # type: ignore[arg-type]
+    elif isinstance(value, (_protocols.DLPackCompatible, _protocols.ArrayCompatible)):
+        return _core.Tensor(value, dtype=dtype, name=name, doc_string=doc_string)
+
+    # Plain (numerical) Python object. Determine the numpy dtype and use np.array to construct the tensor
+    if dtype is not None:
+        if not isinstance(dtype, _enums.DataType):
+            raise TypeError(f"dtype must be an instance of DataType. dtype={dtype}")
+        numpy_dtype = dtype.numpy()
+    elif isinstance(value, Sequence) and not value:
+        raise ValueError("dtype must be specified when value is an empty sequence.")
+    elif isinstance(value, int) and not isinstance(value, bool):
+        # Specify int64 for ints because on Windows this may be int32
+        numpy_dtype = np.dtype(np.int64)
+    elif isinstance(value, float):
+        # If the value is a single float, we use np.float32 as the default dtype
+        numpy_dtype = np.dtype(np.float32)
+    elif isinstance(value, Sequence) and value:
+        if all((isinstance(elem, int) and not isinstance(elem, bool)) for elem in value):
+            numpy_dtype = np.dtype(np.int64)
+        elif all(isinstance(elem, float) for elem in value):
+            # If the value is a sequence of floats, we use np.float32 as the default dtype
+            numpy_dtype = np.dtype(np.float32)
+        else:
+            numpy_dtype = None
+    else:
+        numpy_dtype = None
+
+    array = np.array(value, dtype=numpy_dtype)
+
+    # Handle string tensors by encoding them
+    if isinstance(value, str) or (
+        isinstance(value, Sequence) and value and all(isinstance(elem, str) for elem in value)
+    ):
+        array = np.strings.encode(array, encoding="utf-8")
+        return _core.StringTensor(
+            array,
+            shape=_core.Shape(array.shape),
+            name=name,
+            doc_string=doc_string,
+        )
+
+    return _core.Tensor(
+        array,
+        dtype=dtype,
+        shape=_core.Shape(array.shape),
+        name=name,
+        doc_string=doc_string,
+    )
+
+
+def node(
+    op_type: str,
+    inputs: Sequence[ir.Value | None],
+    attributes: Mapping[str, _convenience.SupportedAttrTypes] | None = None,
+    *,
+    domain: str = "",
+    overload: str = "",
+    num_outputs: int | None = None,
+    outputs: Sequence[ir.Value] | None = None,
+    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.Node:
+    """Create an :class:`ir.Node`.
+
+    This is a convenience constructor for creating a Node that supports Python
+    objects as attributes.
+
+    Example::
+
+        >>> import onnx_ir as ir
+        >>> input_a = ir.Input("A", shape=ir.Shape([1, 2]), type=ir.TensorType(ir.DataType.INT32))
+        >>> input_b = ir.Input("B", shape=ir.Shape([1, 2]), type=ir.TensorType(ir.DataType.INT32))
+        >>> node = ir.node(
+        ...     "SomeOp",
+        ...     inputs=[input_a, input_b],
+        ...     attributes={"alpha": 1.0, "some_list": [1, 2, 3]},
+        ...     domain="some.domain",
+        ...     name="node_name"
+        ... )
+        >>> node.op_type
+        'SomeOp'
+
+    Args:
+        op_type: The name of the operator.
+        inputs: The input values. When an input is None, it is an empty input.
+        attributes: The attributes. RefAttr can be used only when the node is defined in a Function.
+        overload: The overload name when the node is invoking a function.
+        domain: The domain of the operator. For onnx operators, this is an empty string.
+        num_outputs: The number of outputs of the node. If not specified, the number is 1.
+        outputs: The output values. If None, the outputs are created during initialization.
+        version: The version of the operator. If None, the version is unspecified and will follow that of the graph.
+        graph: The graph that the node belongs to. If None, the node is not added to any graph.
+            A `Node` must belong to zero or one graph.
+        name: The name of the node. If None, the node is anonymous.
+        doc_string: The documentation string.
+        metadata_props: The metadata properties.
+
+    Returns:
+        A node with the given op_type and inputs.
+    """
+    if attributes is None:
+        attrs: Sequence[ir.Attr] = ()
+    else:
+        attrs = _convenience.convert_attributes(attributes)
+    return _core.Node(
+        domain=domain,
+        op_type=op_type,
+        inputs=inputs,
+        attributes=attrs,
+        overload=overload,
+        num_outputs=num_outputs,
+        outputs=outputs,
+        version=version,
+        graph=graph,
+        name=name,
+        doc_string=doc_string,
+        metadata_props=metadata_props,
+    )