onnx-ir 0.1.15__py3-none-any.whl

onnx_ir/_display.py

@@ -0,0 +1,54 @@
+# 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
+importing unnecessary dependencies.
+"""
+# pylint: disable=import-outside-toplevel
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def require_rich() -> Any:
+    """Raise an ImportError if rich is not installed."""
+    try:
+        import rich
+    except ImportError:
+        return None
+    return rich
+
+
+class PrettyPrintable:
+    def display(self, *, page: bool = False) -> None:
+        """Pretty print the object.
+
+        Args:
+            page: Whether to page the output.
+        """
+        rich = require_rich()
+        text = str(self)
+
+        if rich is None:
+            print(text)
+            # Color print this message
+            print(
+                f"\n\n\u001b[36mTip: Install the rich library with 'pip install rich' to pretty print this {self.__class__.__name__}.\u001b[0m"
+            )
+            return
+
+        import rich.markup
+
+        # Escape text to display `[...]` correctly
+        text = rich.markup.escape(text)
+
+        if page:
+            import rich.console
+
+            console = rich.console.Console()
+            with console.pager():
+                console.print(text)
+        else:
+            rich.print(text)

onnx_ir/_enums.py

@@ -0,0 +1,474 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""ONNX IR enums that matches the ONNX spec."""
+
+from __future__ import annotations
+
+import enum
+from typing import Any
+
+import ml_dtypes
+import numpy as np
+
+
+class AttributeType(enum.IntEnum):
+    """Enum for the types of ONNX attributes."""
+
+    UNDEFINED = 0
+    FLOAT = 1
+    INT = 2
+    STRING = 3
+    TENSOR = 4
+    GRAPH = 5
+    FLOATS = 6
+    INTS = 7
+    STRINGS = 8
+    TENSORS = 9
+    GRAPHS = 10
+    SPARSE_TENSOR = 11
+    SPARSE_TENSORS = 12
+    TYPE_PROTO = 13
+    TYPE_PROTOS = 14
+
+    def __repr__(self) -> str:
+        return self.name
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+
+class DataType(enum.IntEnum):
+    """Enum for the data types of ONNX tensors, defined in ``onnx.TensorProto``."""
+
+    # NOTE: Naming: It is tempting to use shorter and more modern names like f32, i64,
+    # but we should stick to the names used in the ONNX spec for consistency.
+    UNDEFINED = 0
+    FLOAT = 1
+    UINT8 = 2
+    INT8 = 3
+    UINT16 = 4
+    INT16 = 5
+    INT32 = 6
+    INT64 = 7
+    STRING = 8
+    BOOL = 9
+    FLOAT16 = 10
+    DOUBLE = 11
+    UINT32 = 12
+    UINT64 = 13
+    COMPLEX64 = 14
+    COMPLEX128 = 15
+    BFLOAT16 = 16
+    FLOAT8E4M3FN = 17
+    FLOAT8E4M3FNUZ = 18
+    FLOAT8E5M2 = 19
+    FLOAT8E5M2FNUZ = 20
+    UINT4 = 21
+    INT4 = 22
+    FLOAT4E2M1 = 23
+    FLOAT8E8M0 = 24
+    UINT2 = 25
+    INT2 = 26
+
+    @classmethod
+    def from_numpy(cls, dtype: np.dtype) -> DataType:
+        """Returns the ONNX data type for the numpy dtype.
+
+        Raises:
+            TypeError: If the data type is not supported by ONNX.
+        """
+        if dtype in _NP_TYPE_TO_DATA_TYPE:
+            return cls(_NP_TYPE_TO_DATA_TYPE[dtype])
+
+        if np.issubdtype(dtype, np.str_) or np.issubdtype(dtype, np.bytes_):
+            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
+        # TODO(#137): Remove this when ONNX 1.19 is the minimum requirement
+        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
+            if dtype.names == ("int2",):
+                return DataType.INT2
+            if dtype.names == ("uint2",):
+                return DataType.UINT2
+        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:
+        """Returns the size of the data type in bytes."""
+        return self.bitwidth / 8
+
+    @property
+    def bitwidth(self) -> int:
+        """Returns the bit width of the data type.
+
+        .. versionadded:: 0.1.2
+
+        Raises:
+            TypeError: If the data type is not supported.
+        """
+        if self not in _BITWIDTH_MAP:
+            raise TypeError(f"Bitwidth not available for ONNX data type: {self}")
+        return _BITWIDTH_MAP[self]
+
+    @property
+    def exponent_bitwidth(self) -> int:
+        """Returns the bit width of the exponent for floating-point types.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not supported.
+        """
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).nexp
+
+        raise TypeError(f"Exponent not available for ONNX data type: {self}")
+
+    @property
+    def mantissa_bitwidth(self) -> int:
+        """Returns the bit width of the mantissa for floating-point types.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not supported.
+        """
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).nmant
+
+        raise TypeError(f"Mantissa not available for ONNX data type: {self}")
+
+    @property
+    def eps(self) -> int | np.floating[Any]:
+        """Returns the difference between 1.0 and the next smallest representable float larger than 1.0 for the ONNX data type.
+
+        Returns 1 for integers.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not a numeric data type.
+        """
+        if self.is_integer():
+            return 1
+
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).eps
+
+        raise TypeError(f"Eps not available for ONNX data type: {self}")
+
+    @property
+    def tiny(self) -> int | np.floating[Any]:
+        """Returns the smallest positive non-zero value for the ONNX data type.
+
+        Returns 1 for integers.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not a numeric data type.
+        """
+        if self.is_integer():
+            return 1
+
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).tiny
+
+        raise TypeError(f"Tiny not available for ONNX data type: {self}")
+
+    @property
+    def min(self) -> int | np.floating[Any]:
+        """Returns the minimum representable value for the ONNX data type.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not a numeric data type.
+        """
+        if self.is_integer():
+            return ml_dtypes.iinfo(self.numpy()).min
+
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).min
+
+        raise TypeError(f"Minimum not available for ONNX data type: {self}")
+
+    @property
+    def max(self) -> int | np.floating[Any]:
+        """Returns the maximum representable value for the ONNX data type.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not a numeric data type.
+        """
+        if self.is_integer():
+            return ml_dtypes.iinfo(self.numpy()).max
+
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).max
+
+        raise TypeError(f"Maximum not available for ONNX data type: {self}")
+
+    @property
+    def precision(self) -> int:
+        """Returns the precision for the ONNX dtype if supported.
+
+        For floats returns the approximate number of decimal digits to which
+        this kind of float is precise. Returns 0 for integers.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not a numeric data type.
+        """
+        if self.is_integer():
+            return 0
+
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).precision
+
+        raise TypeError(f"Precision not available for ONNX data type: {self}")
+
+    @property
+    def resolution(self) -> int | np.floating[Any]:
+        """Returns the resolution for the ONNX dtype if supported.
+
+        Returns the approximate decimal resolution of this type, i.e.,
+         10**-precision. Returns 1 for integers.
+
+        .. versionadded:: 0.1.8
+
+        Raises:
+            TypeError: If the data type is not a numeric data type.
+        """
+        if self.is_integer():
+            return 1
+
+        if self.is_floating_point():
+            return ml_dtypes.finfo(self.numpy()).resolution
+
+        raise TypeError(f"Resolution not available for ONNX data type: {self}")
+
+    def numpy(self) -> np.dtype:
+        """Returns the numpy dtype for the ONNX data type.
+
+        Raises:
+            TypeError: If the data type is not supported by numpy.
+        """
+        if self not in _DATA_TYPE_TO_NP_TYPE:
+            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,
+            DataType.FLOAT8E8M0,
+        }
+
+    def is_integer(self) -> bool:
+        """Returns True if the data type is an integer.
+
+        .. versionadded:: 0.1.4
+        """
+        return self in {
+            DataType.UINT8,
+            DataType.INT8,
+            DataType.UINT16,
+            DataType.INT16,
+            DataType.INT32,
+            DataType.INT64,
+            DataType.UINT32,
+            DataType.UINT64,
+            DataType.UINT4,
+            DataType.INT4,
+            DataType.INT2,
+            DataType.UINT2,
+        }
+
+    def is_signed(self) -> bool:
+        """Returns True if the data type is a signed type.
+
+        .. versionadded:: 0.1.4
+        """
+        return self in {
+            DataType.FLOAT,
+            DataType.INT8,
+            DataType.INT16,
+            DataType.INT32,
+            DataType.INT64,
+            DataType.FLOAT16,
+            DataType.DOUBLE,
+            DataType.COMPLEX64,
+            DataType.COMPLEX128,
+            DataType.BFLOAT16,
+            DataType.FLOAT8E4M3FN,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+            DataType.INT4,
+            DataType.FLOAT4E2M1,
+            DataType.FLOAT8E8M0,
+            DataType.INT2,
+        }
+
+    def is_string(self) -> bool:
+        """Returns True if the data type is a string type.
+
+        .. versionadded:: 0.1.8
+        """
+        return self == DataType.STRING
+
+    def __repr__(self) -> str:
+        return self.name
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+
+_BITWIDTH_MAP = {
+    DataType.FLOAT: 32,
+    DataType.UINT8: 8,
+    DataType.INT8: 8,
+    DataType.UINT16: 16,
+    DataType.INT16: 16,
+    DataType.INT32: 32,
+    DataType.INT64: 64,
+    DataType.BOOL: 8,
+    DataType.FLOAT16: 16,
+    DataType.DOUBLE: 64,
+    DataType.UINT32: 32,
+    DataType.UINT64: 64,
+    DataType.COMPLEX64: 64,  # 2 * 32
+    DataType.COMPLEX128: 128,  # 2 * 64
+    DataType.BFLOAT16: 16,
+    DataType.FLOAT8E4M3FN: 8,
+    DataType.FLOAT8E4M3FNUZ: 8,
+    DataType.FLOAT8E5M2: 8,
+    DataType.FLOAT8E5M2FNUZ: 8,
+    DataType.UINT4: 4,
+    DataType.INT4: 4,
+    DataType.FLOAT4E2M1: 4,
+    DataType.FLOAT8E8M0: 8,
+    DataType.INT2: 2,
+    DataType.UINT2: 2,
+}
+
+
+# We use ml_dtypes to support dtypes that are not in numpy.
+_NP_TYPE_TO_DATA_TYPE = {
+    np.dtype("bool"): DataType.BOOL,
+    np.dtype("complex128"): DataType.COMPLEX128,
+    np.dtype("complex64"): DataType.COMPLEX64,
+    np.dtype("float16"): DataType.FLOAT16,
+    np.dtype("float32"): DataType.FLOAT,
+    np.dtype("float64"): DataType.DOUBLE,
+    np.dtype("int16"): DataType.INT16,
+    np.dtype("int32"): DataType.INT32,
+    np.dtype("int64"): DataType.INT64,
+    np.dtype("int8"): DataType.INT8,
+    np.dtype("object"): DataType.STRING,
+    np.dtype("uint16"): DataType.UINT16,
+    np.dtype("uint32"): DataType.UINT32,
+    np.dtype("uint64"): DataType.UINT64,
+    np.dtype("uint8"): DataType.UINT8,
+    np.dtype(ml_dtypes.bfloat16): DataType.BFLOAT16,
+    np.dtype(ml_dtypes.float8_e4m3fn): DataType.FLOAT8E4M3FN,
+    np.dtype(ml_dtypes.float8_e4m3fnuz): DataType.FLOAT8E4M3FNUZ,
+    np.dtype(ml_dtypes.float8_e5m2): DataType.FLOAT8E5M2,
+    np.dtype(ml_dtypes.float8_e5m2fnuz): DataType.FLOAT8E5M2FNUZ,
+    np.dtype(ml_dtypes.float8_e8m0fnu): DataType.FLOAT8E8M0,
+    np.dtype(ml_dtypes.int4): DataType.INT4,
+    np.dtype(ml_dtypes.uint4): DataType.UINT4,
+    np.dtype(ml_dtypes.float4_e2m1fn): DataType.FLOAT4E2M1,
+    np.dtype(ml_dtypes.int2): DataType.INT2,
+    np.dtype(ml_dtypes.uint2): DataType.UINT2,
+}
+
+# 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.FLOAT8E8M0: "f8e8m0",
+    DataType.FLOAT4E2M1: "f4e2m1",
+    DataType.COMPLEX64: "c64",
+    DataType.COMPLEX128: "c128",
+    DataType.INT2: "i2",
+    DataType.INT4: "i4",
+    DataType.INT8: "i8",
+    DataType.INT16: "i16",
+    DataType.INT32: "i32",
+    DataType.INT64: "i64",
+    DataType.BOOL: "b8",
+    DataType.UINT2: "u2",
+    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

@@ -0,0 +1,23 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for comparing IR graphs."""
+
+from __future__ import annotations
+
+from onnx_ir import _core
+
+# NOTE(justinchuby): We need to ensure a graph has valid inputs and outputs
+# NOTE(justinchuby): A graph may be specified with a set of inputs and outputs
+
+
+def topologically_equal(graph1: _core.Graph, graph2: _core.Graph) -> bool:
+    """Return true if the two graphs are topologically equivalent, without considering initializers.
+
+    Args:
+        graph1: The first graph to compare.
+        graph2: The second graph to compare.
+
+    Returns:
+        True if the graphs are equal, False otherwise.
+    """
+    raise NotImplementedError()