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

onnx_ir/_display.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
 """Internal utilities for displaying the intermediate representation of a model.
 
 NOTE: All third-party imports should be scoped and imported only when used to avoid

onnx_ir/_enums.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
 """ONNX IR enums that matches the ONNX spec."""
 
 from __future__ import annotations
@@ -64,6 +64,7 @@ class DataType(enum.IntEnum):
     FLOAT8E5M2FNUZ = 20
     UINT4 = 21
     INT4 = 22
+    FLOAT4E2M1 = 23
 
     @classmethod
     def from_numpy(cls, dtype: np.dtype) -> DataType:
@@ -72,9 +73,43 @@ class DataType(enum.IntEnum):
         Raises:
             TypeError: If the data type is not supported by ONNX.
         """
-        if dtype not in _NP_TYPE_TO_DATA_TYPE:
-            raise TypeError(f"Unsupported numpy data type: {dtype}")
-        return cls(_NP_TYPE_TO_DATA_TYPE[dtype])
+        if dtype in _NP_TYPE_TO_DATA_TYPE:
+            return cls(_NP_TYPE_TO_DATA_TYPE[dtype])
+
+        if np.issubdtype(dtype, np.str_):
+            return DataType.STRING
+
+        # Special cases for handling custom dtypes defined in ONNX (as of onnx 1.18)
+        # Ref: https://github.com/onnx/onnx/blob/2d42b6a60a52e925e57c422593e88cc51890f58a/onnx/_custom_element_types.py
+        if hasattr(dtype, "names"):
+            if dtype.names == ("bfloat16",):
+                return DataType.BFLOAT16
+            if dtype.names == ("e4m3fn",):
+                return DataType.FLOAT8E4M3FN
+            if dtype.names == ("e4m3fnuz",):
+                return DataType.FLOAT8E4M3FNUZ
+            if dtype.names == ("e5m2",):
+                return DataType.FLOAT8E5M2
+            if dtype.names == ("e5m2fnuz",):
+                return DataType.FLOAT8E5M2FNUZ
+            if dtype.names == ("uint4",):
+                return DataType.UINT4
+            if dtype.names == ("int4",):
+                return DataType.INT4
+            if dtype.names == ("float4e2m1",):
+                return DataType.FLOAT4E2M1
+        raise TypeError(f"Unsupported numpy data type: {dtype}")
+
+    @classmethod
+    def from_short_name(cls, short_name: str) -> DataType:
+        """Returns the ONNX data type for the short name.
+
+        Raises:
+            TypeError: If the short name is not available for the data type.
+        """
+        if short_name not in _SHORT_NAME_TO_DATA_TYPE:
+            raise TypeError(f"Unknown short name: {short_name}")
+        return cls(_SHORT_NAME_TO_DATA_TYPE[short_name])
 
     @property
     def itemsize(self) -> float:
@@ -91,6 +126,36 @@ class DataType(enum.IntEnum):
             raise TypeError(f"Numpy does not support ONNX data type: {self}")
         return _DATA_TYPE_TO_NP_TYPE[self]
 
+    def short_name(self) -> str:
+        """Returns the short name of the data type.
+
+        The short name is a string that is used to represent the data type in a more
+        compact form. For example, the short name for `DataType.FLOAT` is "f32".
+        To get the corresponding data type back, call ``from_short_name`` on a string.
+
+        Naming reference: https://github.com/pytorch/pytorch/blob/4bead7b85ea4160243c74109e0ce9bb80686d016/torch/utils/_dtype_abbrs.py
+
+        Raises:
+            TypeError: If the short name is not available for the data type.
+        """
+        if self not in _DATA_TYPE_TO_SHORT_NAME:
+            raise TypeError(f"Short name not available for ONNX data type: {self}")
+        return _DATA_TYPE_TO_SHORT_NAME[self]
+
+    def is_floating_point(self) -> bool:
+        """Returns True if the data type is a floating point type."""
+        return self in {
+            DataType.FLOAT,
+            DataType.FLOAT16,
+            DataType.DOUBLE,
+            DataType.BFLOAT16,
+            DataType.FLOAT8E4M3FN,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+            DataType.FLOAT4E2M1,
+        }
+
     def __repr__(self) -> str:
         return self.name
 
@@ -121,6 +186,7 @@ _ITEMSIZE_MAP = {
     DataType.FLOAT8E5M2FNUZ: 1,
     DataType.UINT4: 0.5,
     DataType.INT4: 0.5,
+    DataType.FLOAT4E2M1: 0.5,
 }
 
 
@@ -150,5 +216,41 @@ _NP_TYPE_TO_DATA_TYPE = {
     np.dtype(ml_dtypes.uint4): DataType.UINT4,
 }
 
+# TODO(after min req for ml_dtypes>=0.5): Move this inside _NP_TYPE_TO_DATA_TYPE
+_NP_TYPE_TO_DATA_TYPE.update(
+    {np.dtype(ml_dtypes.float4_e2m1fn): DataType.FLOAT4E2M1}
+    if hasattr(ml_dtypes, "float4_e2m1fn")
+    else {}
+)
+
 # ONNX DataType to Numpy dtype.
 _DATA_TYPE_TO_NP_TYPE = {v: k for k, v in _NP_TYPE_TO_DATA_TYPE.items()}
+
+_DATA_TYPE_TO_SHORT_NAME = {
+    DataType.UNDEFINED: "undefined",
+    DataType.BFLOAT16: "bf16",
+    DataType.DOUBLE: "f64",
+    DataType.FLOAT: "f32",
+    DataType.FLOAT16: "f16",
+    DataType.FLOAT8E4M3FN: "f8e4m3fn",
+    DataType.FLOAT8E5M2: "f8e5m2",
+    DataType.FLOAT8E4M3FNUZ: "f8e4m3fnuz",
+    DataType.FLOAT8E5M2FNUZ: "f8e5m2fnuz",
+    DataType.FLOAT4E2M1: "f4e2m1",
+    DataType.COMPLEX64: "c64",
+    DataType.COMPLEX128: "c128",
+    DataType.INT4: "i4",
+    DataType.INT8: "i8",
+    DataType.INT16: "i16",
+    DataType.INT32: "i32",
+    DataType.INT64: "i64",
+    DataType.BOOL: "b8",
+    DataType.UINT4: "u4",
+    DataType.UINT8: "u8",
+    DataType.UINT16: "u16",
+    DataType.UINT32: "u32",
+    DataType.UINT64: "u64",
+    DataType.STRING: "s",
+}
+
+_SHORT_NAME_TO_DATA_TYPE = {v: k for k, v in _DATA_TYPE_TO_SHORT_NAME.items()}

onnx_ir/_graph_comparison.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
 """Utilities for comparing IR graphs."""
 
 from __future__ import annotations

onnx_ir/_graph_containers.py

@@ -0,0 +1,268 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Tracked containers for graph."""
+
+# pylint: disable=protected-access
+
+from __future__ import annotations
+
+__all__ = [
+    "GraphInputs",
+    "GraphOutputs",
+]
+
+import collections
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, SupportsIndex
+
+import onnx_ir
+
+if TYPE_CHECKING:
+    from onnx_ir import _core
+
+
+class _GraphIO(collections.UserList["_core.Value"]):
+    """The inputs and outputs of a Graph."""
+
+    def __init__(self, graph: _core.Graph, initlist=None):
+        self._graph = graph
+        # Use a ref counter to track the number of references to each value
+        # in the input/output list. This is used to determine when to unset the graph
+        # reference in the value.
+        # Even though a duplicated value is invalid in inputs and not recommended in outputs,
+        # it is still possible to have duplicated inputs/outputs in an ONNX graph so we
+        # need to properly handle this case and maintain the graph reference properly.
+        self._ref_counter: collections.Counter[_core.Value] = collections.Counter()
+        if initlist is not None:
+            initlist = tuple(initlist)  # Create a copy in case initlist is a generator
+            for value in initlist:
+                self._set_graph(value)
+        super().__init__(initlist)
+        self._check_invariance()
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        raise NotImplementedError
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        raise NotImplementedError
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        raise NotImplementedError
+
+    def append(self, item: _core.Value) -> None:
+        """Add a new input to the graph."""
+        # Perform checks first in _set_graph before modifying the data structure
+        self._set_graph(item)
+        super().append(item)
+        self._check_invariance()
+
+    def extend(self, other) -> None:
+        """Extend the list of inputs or outputs."""
+        other = tuple(other)
+        for item in other:
+            self._set_graph(item)
+        super().extend(other)
+
+    def insert(self, i: int, item: _core.Value) -> None:
+        """Insert an input/output to the graph."""
+        super().insert(i, item)
+        self._set_graph(item)
+        self._check_invariance()
+
+    def pop(self, i: int = -1) -> _core.Value:
+        """Remove an input/output from the graph."""
+        value = super().pop(i)
+        self._maybe_unset_graph(value)
+        self._check_invariance()
+        return value
+
+    def remove(self, item: _core.Value) -> None:
+        """Remove an input/output from the graph."""
+        super().remove(item)
+        self._maybe_unset_graph(item)
+        self._check_invariance()
+
+    def clear(self) -> None:
+        """Clear the list."""
+        for value in self.data:
+            self._maybe_unset_graph(value)
+        super().clear()
+
+    def copy(self) -> list[_core.Value]:
+        """Return a shallow copy of the list."""
+        # This is a shallow copy, so the values are not copied, just the references
+        return self.data.copy()
+
+    def __setitem__(self, i, item) -> None:
+        """Replace an input/output to the node."""
+        if isinstance(item, Iterable) and isinstance(i, slice):
+            # Modify a slice of the list
+            for value in self.data[i]:
+                self._maybe_unset_graph(value)
+            for value in item:
+                self._set_graph(value)
+            super().__setitem__(i, item)
+            self._check_invariance()
+            return
+        elif isinstance(i, SupportsIndex):
+            # Replace a single item
+            self._maybe_unset_graph(self.data[i])
+            self._set_graph(item)
+            super().__setitem__(i, item)
+            self._check_invariance()
+            return
+
+        raise TypeError(f"Invalid types for __setitem__: {type(i)} and {type(item)}")
+
+    def __getitem__(self, i):
+        """Get an input/output from the graph."""
+        return self.data[i]
+
+    def _unimplemented(self, *_args, **_kwargs):
+        """Unimplemented method."""
+        raise RuntimeError("Method is not supported")
+
+    __add__ = _unimplemented
+    __radd__ = _unimplemented
+    __iadd__ = _unimplemented
+    __mul__ = _unimplemented
+    __rmul__ = _unimplemented
+
+
+class GraphInputs(_GraphIO):
+    """The inputs of a Graph."""
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        if not onnx_ir.DEBUG:
+            return
+        for value in self.data:
+            if value._graph is self._graph:
+                continue
+            raise ValueError(
+                f"Invariance error: Value '{value}' is not an input of the graph: {self._graph!r}"
+            )
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph is not None and value._graph is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already owned by a different graph. Please remove the value from the previous graph first"
+            )
+        self._ref_counter[value] += 1
+        value._is_graph_input = True
+        value._graph = self._graph
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        assert value._graph is self._graph, "Bug: value does not belong to the graph"
+        self._ref_counter[value] -= 1
+        if self._ref_counter[value] > 0:
+            # The value is still used by another graph input
+            return
+        value._is_graph_input = False
+        if value._owned_by_graph():
+            # Keep the graph reference if the value is still an input or an initializer
+            return
+        value._graph = None
+
+
+class GraphOutputs(_GraphIO):
+    """The outputs of a Graph."""
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        if not onnx_ir.DEBUG:
+            return
+        for value in self.data:
+            if value._graph is self._graph:
+                continue
+            raise ValueError(
+                f"Invariance error: Value '{value}' is not an output of the graph: {self._graph!r}"
+            )
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph is not None and value._graph is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an output of a different graph. Please remove the value from the previous graph first"
+            )
+        self._ref_counter[value] += 1
+        value._is_graph_output = True
+        value._graph = self._graph
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        assert value._graph is self._graph, "Bug: value does not belong to the graph"
+        self._ref_counter[value] -= 1
+        if self._ref_counter[value] > 0:
+            # The value is still used by another graph input
+            return
+        value._is_graph_output = False
+        if value._owned_by_graph():
+            # Keep the graph reference if the value is still an input or an initializer
+            return
+        value._graph = None
+
+
+class GraphInitializers(collections.UserDict[str, "_core.Value"]):
+    """The initializers of a Graph."""
+
+    def __init__(self, graph: _core.Graph, dict=None, /, **kwargs):
+        # Perform checks first in _set_graph before modifying the data structure with super().__init__()
+        data = {}
+        if dict is not None:
+            data.update(dict)
+        if kwargs:
+            data.update(kwargs)
+        self._graph = graph
+        for value in data.values():
+            self._set_graph(value)
+
+        super().__init__(data)
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph is not None and value._graph is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an initializer of a different graph. Please remove the value from the previous graph first"
+            )
+        value._is_initializer = True
+        value._graph = self._graph
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        assert value._graph is self._graph, "Bug: value does not belong to the graph"
+        value._is_initializer = False
+        if value._owned_by_graph():
+            # Keep the graph reference if the value is still an input or an initializer
+            return
+        value._graph = None
+
+    def __setitem__(self, key: str, value: _core.Value) -> None:
+        """Set an initializer for the graph."""
+        if key != value.name:
+            raise ValueError(
+                f"Key '{key}' does not match the name of the value '{value.name}'"
+            )
+        if not isinstance(key, str):
+            raise TypeError(f"Key must be a string, not {type(key)}")
+        if key in self.data:
+            # If the key already exists, unset the old value
+            old_value = self.data[key]
+            self._maybe_unset_graph(old_value)
+        # Must call _set_graph before super().__setitem__ so that when there is an error,
+        # the dictionary is not modified
+        self._set_graph(value)
+        super().__setitem__(key, value)
+
+    def __delitem__(self, key: str) -> None:
+        """Delete an initializer from the graph."""
+        value = self.data[key]
+        # Must call _maybe_unset_graph before super().__delitem__ so that when there is an error,
+        # the dictionary is not modified
+        self._maybe_unset_graph(value)
+        super().__delitem__(key)

onnx_ir/_io.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
 """Load and save ONNX models."""
 
 from __future__ import annotations
@@ -10,7 +10,9 @@ import os
 
 import onnx
 
-from onnx_ir import _core, _external_data, serde
+from onnx_ir import _core, serde
+from onnx_ir import external_data as _external_data
+from onnx_ir._polyfill import zip
 
 
 def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
@@ -35,16 +37,61 @@ def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
     return model
 
 
-def save(model: _core.Model, path: str | os.PathLike, format: str | None = None) -> None:
+def save(
+    model: _core.Model,
+    path: str | os.PathLike,
+    format: str | None = None,
+    external_data: str | os.PathLike | None = None,
+    size_threshold_bytes: int = 256,
+) -> None:
     """Save an ONNX model to a file.
 
+    The model remains unchanged after the call. If any existing external tensor
+    references the provided ``external_data`` path, it will be invalidated
+    after the external data is overwritten. To obtain a valid model, use :func:`load`
+    to load the newly saved model, or provide a different external data path that
+    is not currently referenced by any tensors in the model.
+
     Args:
         model: The model to save.
-        path: The path to save the model to.
-        format: The format of the file (e.g. protobuf, textproto, json, etc.).
+        path: The path to save the model to. E.g. "model.onnx".
+        format: The format of the file (e.g. ``protobuf``, ``textproto``, ``json``, etc.).
             If None, the format is inferred from the file extension.
+        external_data: The relative path to save external data to. When specified,
+            all initializers in the model will be converted to external data and
+            saved to the specified directory. If None, all tensors will be saved unmodified.
+            That is, if a tensor in the model is already external, it will be saved
+            with the same external information; if the tensor is not external,
+            it will be serialized in the ONNX Proto message.
+        size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
+            Effective only when ``external_data`` is set.
+
+    Raises:
+        ValueError: If the external data path is an absolute path.
     """
-    proto = serde.serialize_model(model)
-    onnx.save(proto, path, format=format)
-    # TODO(justinchuby): Handle external data when the relative path has changed
-    # TODO(justinchuby): Handle off loading external data to disk when saving
+    if external_data is not None:
+        if os.path.isabs(external_data):
+            raise ValueError(
+                f"The external data path must be relative to the ONNX file path, not '{external_data}'."
+            )
+        base_dir = os.path.dirname(path)
+
+        # Store the original initializer values so they can be restored if modify_model=False
+        initializer_values = tuple(model.graph.initializers.values())
+        tensors = [v.const_value for v in initializer_values]
+
+        try:
+            model = _external_data.unload_from_model(
+                model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes
+            )
+            proto = serde.serialize_model(model)
+            onnx.save(proto, path, format=format)
+
+        finally:
+            # Restore the original initializer values so the model is unchanged
+            for initializer, tensor in zip(initializer_values, tensors, strict=True):
+                initializer.const_value = tensor
+
+    else:
+        proto = serde.serialize_model(model)
+        onnx.save(proto, path, format=format)

onnx_ir/_linked_list.py

@@ -1,10 +1,11 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """Mutable list for nodes in a graph with safe mutation properties."""
 
 from __future__ import annotations
 
-from typing import Generic, Iterable, Iterator, Sequence, TypeVar
+from collections.abc import Iterable, Iterator, Sequence
+from typing import Generic, TypeVar, overload
 
 T = TypeVar("T")
 
@@ -131,16 +132,23 @@ class DoublyLinkedSet(Sequence[T], Generic[T]):
             box = box.prev
 
     def __len__(self) -> int:
-        assert self._length == len(
-            self._value_ids_to_boxes
-        ), "Bug in the implementation: length mismatch"
+        assert self._length == len(self._value_ids_to_boxes), (
+            "Bug in the implementation: length mismatch"
+        )
         return self._length
 
-    def __getitem__(self, index: int) -> T:
+    @overload
+    def __getitem__(self, index: int) -> T: ...
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[T]: ...
+
+    def __getitem__(self, index):
         """Get the node at the given index.
 
         Complexity is O(n).
         """
+        if isinstance(index, slice):
+            return tuple(self)[index]
         if index >= self._length or index < -self._length:
             raise IndexError(
                 f"Index out of range: {index} not in range [-{self._length}, {self._length})"

onnx_ir/_metadata.py

@@ -1,11 +1,12 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 """Class for storing metadata about the IR objects."""
 
 from __future__ import annotations
 
 import collections
-from typing import Any, Mapping
+from collections.abc import Mapping
+from typing import Any
 
 
 class MetadataStore(collections.UserDict):

onnx_ir/_name_authority.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
 """Auxiliary class for managing names in the IR."""
 
 from __future__ import annotations

onnx_ir/_polyfill.py

@@ -0,0 +1,26 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Polyfill for Python builtin functions."""
+
+import sys
+from collections.abc import Sequence
+from typing import Any
+
+if sys.version_info >= (3, 10):
+    zip = zip  # pylint: disable=self-assigning-variable
+else:
+    # zip(..., strict=True) was added in Python 3.10
+    # TODO: Remove this polyfill when we drop support for Python 3.9
+    _python_zip = zip
+
+    def zip(a: Sequence[Any], b: Sequence[Any], strict: bool = False):
+        """Polyfill for Python's zip function.
+
+        This is a special version which only supports two Sequence inputs.
+
+        Raises:
+            ValueError: If the iterables have different lengths and strict is True.
+        """
+        if len(a) != len(b) and strict:
+            raise ValueError("zip() argument lengths must be equal")
+        return _python_zip(a, b)