onnx-ir 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

onnx_ir/__init__.py

@@ -78,6 +78,7 @@ __all__ = [
     # Convenience constructors
     "tensor",
     "node",
+    "val",
     # Pass infrastructure
     "passes",
     # IO
@@ -90,7 +91,7 @@ __all__ = [
 import types
 
 from onnx_ir import convenience, external_data, passes, serde, tape, traversal
-from onnx_ir._convenience._constructors import node, tensor
+from onnx_ir._convenience._constructors import node, tensor, val
 from onnx_ir._core import (
     Attr,
     AttrFloat32,
@@ -167,4 +168,4 @@ def __set_module() -> None:
 
 
 __set_module()
-__version__ = "0.1.8"
+__version__ = "0.1.10"

onnx_ir/_convenience/__init__.py

@@ -226,7 +226,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]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', GRAPH, 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', GRAPH, Graph(
             name='graph0',
             inputs=(
         <BLANKLINE>
@@ -235,7 +235,7 @@ def convert_attributes(
         <BLANKLINE>
             ),
             len()=0
-        )), Attr('graphs', GRAPHS, [Graph(
+        )), Attr('graphs', GRAPHS, (Graph(
             name='graph1',
             inputs=(
         <BLANKLINE>
@@ -253,7 +253,7 @@ def convert_attributes(
         <BLANKLINE>
             ),
             len()=0
-        )]), Attr('type_proto', TYPE_PROTO, Tensor(FLOAT)), Attr('type_protos', TYPE_PROTOS, [Tensor(FLOAT), Tensor(FLOAT)])]
+        ))), Attr('type_proto', TYPE_PROTO, Tensor(FLOAT)), Attr('type_protos', TYPE_PROTOS, (Tensor(FLOAT), Tensor(FLOAT)))]
 
     .. important::
         An empty sequence should be created with an explicit type by initializing
@@ -293,7 +293,7 @@ def replace_all_uses_with(
     We want to replace the node A with a new node D::
 
         >>> import onnx_ir as ir
-        >>> input = ir.Input("input")
+        >>> input = ir.val("input")
         >>> node_a = ir.Node("", "A", [input])
         >>> node_b = ir.Node("", "B", node_a.outputs)
         >>> node_c = ir.Node("", "C", node_a.outputs)

onnx_ir/_convenience/_constructors.py

@@ -25,7 +25,7 @@ if typing.TYPE_CHECKING:
 
 def tensor(
     value: npt.ArrayLike | onnx.TensorProto | ir.DLPackCompatible | ir.ArrayCompatible,
-    dtype: _enums.DataType | None = None,
+    dtype: ir.DataType | None = None,
     name: str | None = None,
     doc_string: str | None = None,
 ) -> _protocols.TensorProtocol:
@@ -159,7 +159,7 @@ def node(
     doc_string: str | None = None,
     metadata_props: dict[str, str] | None = None,
 ) -> ir.Node:
-    """Create an :class:`~onnx_ir.Node`.
+    """Create a :class:`~onnx_ir.Node`.
 
     This is a convenience constructor for creating a Node that supports Python
     objects as attributes.
@@ -167,8 +167,8 @@ def node(
     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))
+        >>> input_a = ir.val("A", shape=[1, 2], type=ir.TensorType(ir.DataType.INT32))
+        >>> input_b = ir.val("B", shape=[1, 2], type=ir.TensorType(ir.DataType.INT32))
         >>> node = ir.node(
         ...     "SomeOp",
         ...     inputs=[input_a, input_b],
@@ -215,3 +215,77 @@ def node(
         doc_string=doc_string,
         metadata_props=metadata_props,
     )
+
+
+def val(
+    name: str | None,
+    dtype: ir.DataType | None = None,
+    shape: ir.Shape | Sequence[int | str | None] | None = None,
+    *,
+    type: ir.TypeProtocol | None = None,
+    const_value: ir.TensorProtocol | None = None,
+    metadata_props: dict[str, str] | None = None,
+) -> ir.Value:
+    """Create a :class:`~onnx_ir.Value` with the given name and type.
+
+    This is a convenience constructor for creating a Value that allows you to specify
+    dtype and shape in a more relaxed manner. Whereas to create a Value directly, you
+    need to create a :class:`~onnx_ir.TypeProtocol` and :class:`~onnx_ir.Shape` object
+    first, this function allows you to specify dtype as a :class:`~onnx_ir.DataType`
+    and shape as a sequence of integers or symbolic dimensions.
+
+    Example::
+
+        >>> import onnx_ir as ir
+        >>> t = ir.val("x", ir.DataType.FLOAT, ["N", 42, 3])
+        >>> t.name
+        'x'
+        >>> t.type
+        Tensor(FLOAT)
+        >>> t.shape
+        Shape([SymbolicDim(N), 42, 3])
+
+    .. versionadded:: 0.1.9
+
+    Args:
+        name: The name of the value.
+        dtype: The data type of the TensorType of the value. This is used only when type is None.
+        shape: The shape of the value.
+        type: The type of the value. Only one of dtype and type can be specified.
+        const_value: The constant tensor that initializes the value. Supply this argument
+            when you want to create an initializer. The type and shape can be obtained from the tensor.
+        metadata_props: The metadata properties that will be serialized to the ONNX proto.
+
+    Returns:
+        A Value object.
+    """
+    if const_value is not None:
+        const_tensor_type = _core.TensorType(const_value.dtype)
+        if type is not None and type != const_tensor_type:
+            raise ValueError(
+                f"The type does not match the const_value. type={type} but const_value has type {const_tensor_type}. "
+                "You do not have to specify the type when const_value is provided."
+            )
+        if dtype is not None and dtype != const_value.dtype:
+            raise ValueError(
+                f"The dtype does not match the const_value. dtype={dtype} but const_value has dtype {const_value.dtype}. "
+                "You do not have to specify the dtype when const_value is provided."
+            )
+        if shape is not None and _core.Shape(shape) != const_value.shape:
+            raise ValueError(
+                f"The shape does not match the const_value. shape={shape} but const_value has shape {const_value.shape}. "
+                "You do not have to specify the shape when const_value is provided."
+            )
+        return _core.Value(
+            name=name,
+            type=const_tensor_type,
+            shape=_core.Shape(const_value.shape),  # type: ignore
+            const_value=const_value,
+            metadata_props=metadata_props,
+        )
+
+    if type is None and dtype is not None:
+        type = _core.TensorType(dtype)
+    if shape is not None and not isinstance(shape, _core.Shape):
+        shape = _core.Shape(shape)
+    return _core.Value(name=name, type=type, shape=shape, metadata_props=metadata_props)

onnx_ir/_core.py

@@ -45,7 +45,7 @@ from typing import (
 
 import ml_dtypes
 import numpy as np
-from typing_extensions import TypeIs
+from typing_extensions import TypeIs, deprecated
 
 import onnx_ir
 from onnx_ir import (
@@ -165,6 +165,11 @@ class TensorBase(abc.ABC, _protocols.TensorProtocol, _display.PrettyPrintable):
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the tensor.
+
+        The metadata properties are used to store additional information about the tensor.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -1214,6 +1219,12 @@ class Shape(_protocols.ShapeProtocol, _display.PrettyPrintable):
 
     Use :meth:`get_denotation` and :meth:`set_denotation` to access and modify the denotations.
 
+    .. note::
+        Two shapes can be compared for equality. Be careful when comparing shapes with
+        unknown dimensions (``None``), as they may not be considered semantically equal
+        even if all dimensions are the same. You can use :meth:`has_unknown_dim` to
+        check if a shape has any unknown dimensions.
+
     Example::
 
         >>> import onnx_ir as ir
@@ -1422,6 +1433,29 @@ class Shape(_protocols.ShapeProtocol, _display.PrettyPrintable):
             return not self.is_static()
         return not self.is_static(dim)
 
+    def is_unknown_dim(self, dim: int) -> bool:
+        """Return True if the dimension is unknown (None).
+
+        A dynamic dimension without a symbolic name is considered unknown.
+
+        .. versionadded:: 0.1.10
+
+        Args:
+            dim: The index of the dimension.
+        """
+        dim_obj = self._dims[dim]
+        return isinstance(dim_obj, SymbolicDim) and dim_obj.value is None
+
+    def has_unknown_dim(self) -> bool:
+        """Return True if any dimension is unknown (None).
+
+        You can use :meth:`is_unknown_dim` to check if a specific dimension is unknown.
+
+        .. versionadded:: 0.1.10
+        """
+        # We can use "in" directly because SymbolicDim implements __eq__ with None
+        return None in self._dims
+
 
 def _quoted(string: str) -> str:
     """Return a quoted string.
@@ -2022,9 +2056,18 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
         type: _protocols.TypeProtocol | None = None,
         doc_string: str | None = None,
         const_value: _protocols.TensorProtocol | None = None,
+        metadata_props: dict[str, str] | None = None,
     ) -> None:
         """Initialize a value.
 
+        When assigning a name to the value, the name of the backing `const_value` (Tensor)
+        will also be updated. If the value is an initializer of a graph, the initializers
+        dictionary of the graph will also be updated.
+
+        .. versionchanged:: 0.1.10
+            Assigning a name to the value will also update the graph initializer entry
+            if the value is an initializer of a graph.
+
         Args:
             producer: The node that produces the value.
                 It can be ``None`` when the value is initialized first than its producer.
@@ -2034,11 +2077,12 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
             type: The type of the value.
             doc_string: The documentation string.
             const_value: The constant tensor if the value is constant.
+            metadata_props: Metadata that will be serialized to the ONNX file.
         """
         self._producer: Node | None = producer
         self._index: int | None = index
         self._metadata: _metadata.MetadataStore | None = None
-        self._metadata_props: dict[str, str] | None = None
+        self._metadata_props: dict[str, str] | None = metadata_props
 
         self._name: str | None = name
         self._shape: Shape | None = shape
@@ -2172,7 +2216,23 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
     def name(self, value: str | None) -> None:
         if self._const_value is not None:
             self._const_value.name = value
+        old_name = self._name
         self._name = value
+        if self.is_initializer():
+            if value is None:
+                raise ValueError(
+                    "Initializer value cannot have name set to None. Please pop() the value from initializers first"
+                )
+            # Rename the initializer entry in the graph
+            graph = self._graph
+            assert graph is not None
+            assert old_name is not None
+            if value in graph.initializers and graph.initializers[value] is not self:
+                raise ValueError(
+                    f"Cannot rename initializer to '{value}': an initializer with that name already exists."
+                )
+            graph.initializers.pop(old_name)
+            graph.initializers[value] = self
 
     @property
     def type(self) -> _protocols.TypeProtocol | None:
@@ -2226,9 +2286,16 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
     def const_value(
         self,
     ) -> _protocols.TensorProtocol | None:
-        """A concrete value.
+        """The backing constant tensor for the value.
+
+        If the ``Value`` has a ``const_value`` and is part of a graph initializers
+        dictionary, the value is an initialized value. Its ``const_value``
+        will appear as an ``initializer`` in the GraphProto when serialized.
+
+        If the ``Value`` is not part of a graph initializers dictionary, the ``const_value``
+        field will be ignored during serialization.
 
-        The value can be backed by different raw data types, such as numpy arrays.
+        ``const_value`` can be backed by different raw data types, such as numpy arrays.
         The only guarantee is that it conforms TensorProtocol.
         """
         return self._const_value
@@ -2258,6 +2325,11 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the value.
+
+        The metadata properties are used to store additional information about the value.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -2275,6 +2347,7 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
         return self._is_initializer
 
 
+@deprecated("Input is deprecated since 0.1.9. Use ir.val(...) instead.")
 def Input(
     name: str | None = None,
     shape: Shape | None = None,
@@ -2804,6 +2877,11 @@ class Graph(_protocols.GraphProtocol, Sequence[Node], _display.PrettyPrintable):
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the graph.
+
+        The metadata properties are used to store additional information about the graph.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -3056,6 +3134,11 @@ class Model(_protocols.ModelProtocol, _display.PrettyPrintable):
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the model.
+
+        The metadata properties are used to store additional information about the model.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -3249,6 +3332,11 @@ class Function(_protocols.FunctionProtocol, Sequence[Node], _display.PrettyPrint
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the function.
+
+        The metadata properties are used to store additional information about the function.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         return self._graph.metadata_props
 
     def all_nodes(self) -> Iterator[Node]:
@@ -3397,6 +3485,31 @@ class Attr(
         *,
         doc_string: str | None = None,
     ) -> None:
+        # Quick checks to ensure that INT and FLOAT attributes are stored as int and float,
+        # not np.int32, np.float32, bool, etc.
+        # This also allows errors to be raised at the time of construction instead of later
+        # during serialization.
+        # TODO(justinchuby): Use case matching when we drop support for Python 3.9
+        if value is None:
+            # Value can be None for reference attributes or when it is used as a
+            # placeholder for schemas
+            pass
+        elif type == _enums.AttributeType.INT:
+            value = int(value)
+        elif type == _enums.AttributeType.FLOAT:
+            value = float(value)
+        elif type == _enums.AttributeType.INTS:
+            value = tuple(int(v) for v in value)
+        elif type == _enums.AttributeType.FLOATS:
+            value = tuple(float(v) for v in value)
+        elif type in {
+            _enums.AttributeType.STRINGS,
+            _enums.AttributeType.TENSORS,
+            _enums.AttributeType.GRAPHS,
+            _enums.AttributeType.TYPE_PROTOS,
+        }:
+            value = tuple(value)
+
         self._name = name
         self._type = type
         self._value = value
@@ -3458,7 +3571,7 @@ class Attr(
             return f"@{self.ref_attr_name}"
         if self.type == _enums.AttributeType.GRAPH:
             return textwrap.indent("\n" + str(self.value), " " * 4)
-        return str(self.value)
+        return repr(self.value)
 
     def __repr__(self) -> str:
         if self.is_ref():
@@ -3472,8 +3585,8 @@ class Attr(
             raise TypeError(
                 f"Attribute '{self.name}' is not of type FLOAT. Actual type: {self.type}"
             )
-        # Do not use isinstance check because it may prevent np.float32 etc. from being used
-        return float(self.value)
+        # value is guaranteed to be a float in the constructor
+        return self.value
 
     def as_int(self) -> int:
         """Get the attribute value as an int."""
@@ -3481,8 +3594,8 @@ class Attr(
             raise TypeError(
                 f"Attribute '{self.name}' is not of type INT. Actual type: {self.type}"
             )
-        # Do not use isinstance check because it may prevent np.int32 etc. from being used
-        return int(self.value)
+        # value is guaranteed to be an int in the constructor
+        return self.value
 
     def as_string(self) -> str:
         """Get the attribute value as a string."""
@@ -3490,9 +3603,10 @@ class Attr(
             raise TypeError(
                 f"Attribute '{self.name}' is not of type STRING. Actual type: {self.type}"
             )
-        if not isinstance(self.value, str):
+        value = self.value
+        if not isinstance(value, str):
             raise TypeError(f"Value of attribute '{self!r}' is not a string.")
-        return self.value
+        return value
 
     def as_tensor(self) -> _protocols.TensorProtocol:
         """Get the attribute value as a tensor."""
@@ -3500,9 +3614,10 @@ class Attr(
             raise TypeError(
                 f"Attribute '{self.name}' is not of type TENSOR. Actual type: {self.type}"
             )
-        if not isinstance(self.value, _protocols.TensorProtocol):
+        value = self.value
+        if not isinstance(value, _protocols.TensorProtocol):
             raise TypeError(f"Value of attribute '{self!r}' is not a tensor.")
-        return self.value
+        return value
 
     def as_graph(self) -> Graph:
         """Get the attribute value as a graph."""
@@ -3510,75 +3625,64 @@ class Attr(
             raise TypeError(
                 f"Attribute '{self.name}' is not of type GRAPH. Actual type: {self.type}"
             )
-        if not isinstance(self.value, Graph):
+        value = self.value
+        if not isinstance(value, Graph):
             raise TypeError(f"Value of attribute '{self!r}' is not a graph.")
-        return self.value
+        return value
 
-    def as_floats(self) -> Sequence[float]:
+    def as_floats(self) -> tuple[float, ...]:
         """Get the attribute value as a sequence of floats."""
         if self.type != _enums.AttributeType.FLOATS:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type FLOATS. Actual type: {self.type}"
             )
-        if not isinstance(self.value, Sequence):
-            raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
-        # Do not use isinstance check on elements because it may prevent np.int32 etc. from being used
-        # Create a copy of the list to prevent mutation
-        return [float(v) for v in self.value]
+        # value is guaranteed to be a sequence of float in the constructor
+        return self.value
 
-    def as_ints(self) -> Sequence[int]:
+    def as_ints(self) -> tuple[int, ...]:
         """Get the attribute value as a sequence of ints."""
         if self.type != _enums.AttributeType.INTS:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type INTS. Actual type: {self.type}"
             )
-        if not isinstance(self.value, Sequence):
-            raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
-        # Do not use isinstance check on elements because it may prevent np.int32 etc. from being used
-        # Create a copy of the list to prevent mutation
-        return list(self.value)
+        # value is guaranteed to be a sequence of int in the constructor
+        return self.value
 
-    def as_strings(self) -> Sequence[str]:
+    def as_strings(self) -> tuple[str, ...]:
         """Get the attribute value as a sequence of strings."""
         if self.type != _enums.AttributeType.STRINGS:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type STRINGS. Actual type: {self.type}"
             )
-        if not isinstance(self.value, Sequence):
-            raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
         if onnx_ir.DEBUG:
             if not all(isinstance(x, str) for x in self.value):
                 raise TypeError(f"Value of attribute '{self!r}' is not a Sequence of strings.")
-        # Create a copy of the list to prevent mutation
-        return list(self.value)
+        # value is guaranteed to be a sequence in the constructor
+        return self.value
 
-    def as_tensors(self) -> Sequence[_protocols.TensorProtocol]:
+    def as_tensors(self) -> tuple[_protocols.TensorProtocol, ...]:
         """Get the attribute value as a sequence of tensors."""
         if self.type != _enums.AttributeType.TENSORS:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type TENSORS. Actual type: {self.type}"
             )
-        if not isinstance(self.value, Sequence):
-            raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
         if onnx_ir.DEBUG:
             if not all(isinstance(x, _protocols.TensorProtocol) for x in self.value):
                 raise TypeError(f"Value of attribute '{self!r}' is not a Sequence of tensors.")
-        # Create a copy of the list to prevent mutation
-        return list(self.value)
+        # value is guaranteed to be a sequence in the constructor
+        return tuple(self.value)
 
-    def as_graphs(self) -> Sequence[Graph]:
+    def as_graphs(self) -> tuple[Graph, ...]:
         """Get the attribute value as a sequence of graphs."""
         if self.type != _enums.AttributeType.GRAPHS:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type GRAPHS. Actual type: {self.type}"
             )
-        if not isinstance(self.value, Sequence):
-            raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
         if onnx_ir.DEBUG:
             if not all(isinstance(x, Graph) for x in self.value):
                 raise TypeError(f"Value of attribute '{self!r}' is not a Sequence of graphs.")
-        # Create a copy of the list to prevent mutation
-        return list(self.value)
+        # value is guaranteed to be a sequence in the constructor
+        return tuple(self.value)
 
 
 # NOTE: The following functions are just for convenience
@@ -3605,7 +3709,7 @@ def RefAttr(
     return Attr(name, type, None, ref_attr_name=ref_attr_name, doc_string=doc_string)
 
 
-def AttrFloat32(name: str, value: float, doc_string: str | None = None) -> Attr:
+def AttrFloat32(name: str, value: float | np.floating, doc_string: str | None = None) -> Attr:
     """Create a float attribute."""
     # NOTE: The function name is capitalized to maintain API backward compatibility.
     return Attr(
@@ -3616,7 +3720,7 @@ def AttrFloat32(name: str, value: float, doc_string: str | None = None) -> Attr:
     )
 
 
-def AttrInt64(name: str, value: int, doc_string: str | None = None) -> Attr:
+def AttrInt64(name: str, value: int | np.integer, doc_string: str | None = None) -> Attr:
     """Create an int attribute."""
     # NOTE: The function name is capitalized to maintain API backward compatibility.
     return Attr(

onnx_ir/_enums.py

@@ -357,7 +357,10 @@ class DataType(enum.IntEnum):
         }
 
     def is_string(self) -> bool:
-        """Returns True if the data type is a string type."""
+        """Returns True if the data type is a string type.
+
+        .. versionadded:: 0.1.8
+        """
         return self == DataType.STRING
 
     def __repr__(self) -> str:
@@ -419,15 +422,9 @@ _NP_TYPE_TO_DATA_TYPE = {
     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,
 }
 
-# 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()}
 

onnx_ir/passes/common/constant_manipulation.py

@@ -69,7 +69,12 @@ class LiftConstantsToInitializersPass(ir.passes.InPlacePass):
                 shape=tensor.shape,  # type: ignore[arg-type]
                 type=ir.TensorType(tensor.dtype),
                 const_value=tensor,
+                # Preserve metadata from Constant value into the onnx model
+                metadata_props=node.outputs[0].metadata_props.copy(),
             )
+            # Preserve value meta from the Constant output for intermediate analysis
+            initializer.meta.update(node.outputs[0].meta)
+
             assert node.graph is not None
             node.graph.register_initializer(initializer)
             # Replace the constant node with the initializer

onnx_ir/passes/common/naming.py

@@ -193,14 +193,8 @@ class NameFixPass(ir.passes.InPlacePass):
         if not value.name:
             modified = self._assign_value_name(value, used_value_names, value_counter)
         else:
-            old_name = value.name
             modified = self._fix_duplicate_value_name(value, used_value_names, value_counter)
-            if modified:
-                assert value.graph is not None
-                if value.is_initializer():
-                    value.graph.initializers.pop(old_name)
-                    # Add the initializer back with the new name
-                    value.graph.initializers.add(value)
+            # initializers dictionary is updated automatically when the Value is renamed
 
         # Record the final name for this value
         assert value.name is not None

onnx_ir/serde.py

@@ -709,9 +709,8 @@ def _deserialize_graph(
         annotation.tensor_name: annotation for annotation in proto.quantization_annotation
     }
 
-    # Create values for initializers and inputs
-    initializer_tensors = [deserialize_tensor(tensor) for tensor in proto.initializer]
-    inputs = [_core.Input(info.name) for info in proto.input]
+    # Create values for inputs
+    inputs = [_core.Value(name=info.name) for info in proto.input]
     for info, value in zip(proto.input, inputs):
         deserialize_value_info_proto(info, value)
 
@@ -725,6 +724,11 @@ def _deserialize_graph(
     # Enter the graph scope by pushing the values for this scope to the stack
     scoped_values.append(values)
 
+    # Build the value info dictionary to allow for quick lookup for this graph scope
+    value_info = {info.name: info for info in proto.value_info}
+
+    # Create values for initializers
+    initializer_tensors = [deserialize_tensor(tensor) for tensor in proto.initializer]
     initializer_values = []
     for i, tensor in enumerate(initializer_tensors):
         initializer_name = tensor.name
@@ -750,6 +754,8 @@ def _deserialize_graph(
                 shape=tensor.shape,  # type: ignore[arg-type]
                 const_value=tensor,
             )
+            if initializer_name in value_info:
+                deserialize_value_info_proto(value_info[initializer_name], initializer_value)
             if initializer_value.name in quantization_annotations:
                 _deserialize_quantization_annotation(
                     quantization_annotations[initializer_value.name], initializer_value
@@ -757,9 +763,6 @@ def _deserialize_graph(
             values[initializer_name] = initializer_value
         initializer_values.append(initializer_value)
 
-    # Build the value info dictionary to allow for quick lookup for this graph scope
-    value_info = {info.name: info for info in proto.value_info}
-
     # Declare values for all node outputs from this graph scope. This is necessary
     # to handle the case where a node in a subgraph uses a value that is declared out
     # of order in the outer graph. Declaring the values first allows us to find the
@@ -869,7 +872,7 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
     Returns:
         An IR Function object representing the ONNX function.
     """
-    inputs = [_core.Input(name) for name in proto.input]
+    inputs = [_core.Value(name=name) for name in proto.input]
     values: dict[str, _core.Value] = {v.name: v for v in inputs}  # type: ignore[misc]
     value_info = {info.name: info for info in getattr(proto, "value_info", [])}
 
@@ -1143,7 +1146,19 @@ def _deserialize_attribute(
     if type_ == _enums.AttributeType.FLOAT:
         return _core.AttrFloat32(name, proto.f, doc_string=doc_string)
     if type_ == _enums.AttributeType.STRING:
-        return _core.AttrString(name, proto.s.decode("utf-8"), doc_string=doc_string)
+        try:
+            return _core.AttrString(name, proto.s.decode("utf-8"), doc_string=doc_string)
+        except UnicodeDecodeError:
+            # Even though onnx.ai/onnx/repo-docs/IR.html#attributes requires the attribute
+            # for strings to be utf-8 encoded bytes, custom ops may still store arbitrary data there
+            logger.warning(
+                "Attribute %r contains invalid UTF-8 bytes. ONNX spec requires string attributes "
+                "to be UTF-8 encoded so the model is invalid. We will skip decoding the attribute and "
+                "use the bytes as attribute value",
+                name,
+            )
+            return _core.Attr(name, type_, proto.s, doc_string=doc_string)
+
     if type_ == _enums.AttributeType.INTS:
         return _core.AttrInt64s(name, proto.ints, doc_string=doc_string)
     if type_ == _enums.AttributeType.FLOATS:
@@ -1378,7 +1393,12 @@ def _should_create_value_info_for_value(value: _protocols.ValueProtocol) -> bool
         True if value info should be created for the value.
     """
     # No need to serialize value info if it is not set
-    if value.shape is None and value.type is None:
+    if (
+        value.shape is None
+        and value.type is None
+        and not value.metadata_props
+        and not value.doc_string
+    ):
         return False
     if not value.name:
         logger.debug("Did not serialize '%s' because its name is empty", value)
@@ -1784,16 +1804,26 @@ def _fill_in_value_for_attribute(
 ) -> None:
     if type_ == _enums.AttributeType.INT:
         # value: int
-        # Cast bool to int, for example
-        attribute_proto.i = int(value)
+        attribute_proto.i = value
         attribute_proto.type = onnx.AttributeProto.INT
     elif type_ == _enums.AttributeType.FLOAT:
         # value: float
-        attribute_proto.f = float(value)
+        attribute_proto.f = value
         attribute_proto.type = onnx.AttributeProto.FLOAT
     elif type_ == _enums.AttributeType.STRING:
         # value: str
-        attribute_proto.s = value.encode("utf-8")
+        if type(value) is bytes:
+            # Even though onnx.ai/onnx/repo-docs/IR.html#attributes requires the attribute
+            # for strings to be utf-8 encoded bytes, custom ops may still store arbitrary data there
+            logger.warning(
+                "Value in attribute %r should be a string but is instead bytes. ONNX "
+                "spec requires string attributes to be UTF-8 encoded so the model is invalid. "
+                "We will skip encoding the attribute and use the bytes as attribute value",
+                attribute_proto.name,
+            )
+            attribute_proto.s = value
+        else:
+            attribute_proto.s = value.encode("utf-8")
         attribute_proto.type = onnx.AttributeProto.STRING
     elif type_ == _enums.AttributeType.INTS:
         # value: Sequence[int]
@@ -1945,11 +1975,26 @@ def serialize_type(type_protocol: _protocols.TypeProtocol) -> onnx.TypeProto:
 @_capture_errors(lambda type_proto, from_: repr(from_))
 def serialize_shape_into(type_proto: onnx.TypeProto, from_: _protocols.ShapeProtocol) -> None:
     value_field = type_proto.WhichOneof("value")
+    if value_field is None:
+        # We cannot write the shape because we do not know where to write it
+        logger.warning(
+            # TODO(justinchuby): Show more context about the value when move everything to an object
+            "The value type for shape %s is not known. Please set type for the value. Skipping serialization",
+            from_,
+        )
+        return
     tensor_type = getattr(type_proto, value_field)
     while not isinstance(tensor_type.elem_type, int):
         # Find the leaf type that has the shape field
         type_proto = tensor_type.elem_type
         value_field = type_proto.WhichOneof("value")
+        if value_field is None:
+            logger.warning(
+                # TODO(justinchuby): Show more context about the value when move everything to an object
+                "The value type for shape %s is not known. Please set type for the value. Skipping serialization",
+                from_,
+            )
+            return
         tensor_type = getattr(type_proto, value_field)
     # When from is empty, we still need to set the shape field to an empty list by touching it
     tensor_type.shape.ClearField("dim")
@@ -1970,5 +2015,8 @@ def serialize_dimension_into(
         dim_proto.dim_value = dim
     elif isinstance(dim, (_core.SymbolicDim, _protocols.SymbolicDimProtocol)):
         if dim.value is not None:
-            # TODO(justinchuby): None is probably not a valid value for dim_param
             dim_proto.dim_param = str(dim.value)
+        # NOTE: None is a valid value for symbolic dimension:
+        # A dimension MAY have neither dim_value nor dim_param set. Such a dimension
+        # represents an unknown dimension unrelated to other unknown dimensions.
+        # Here we will just leave the dim_proto empty.

{onnx_ir-0.1.8.dist-info → onnx_ir-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.8
+Version: 0.1.10
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License-Expression: Apache-2.0
@@ -19,7 +19,7 @@ License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: onnx>=1.16
 Requires-Dist: typing_extensions>=4.10
-Requires-Dist: ml_dtypes
+Requires-Dist: ml_dtypes>=0.5.0
 Dynamic: license-file
 
 # <img src="docs/_static/logo-light.png" alt="ONNX IR" width="250"/>

{onnx_ir-0.1.8.dist-info → onnx_ir-0.1.10.dist-info}/RECORD

@@ -1,7 +1,7 @@
-onnx_ir/__init__.py,sha256=kcifjxB2nP9t4WMrxcJriVQ_WmgpOf0DOfUgsngd_PA,3424
-onnx_ir/_core.py,sha256=EEQyNzWkz_ukCV_9Dc6AUIVAmPErhVjntr4xwPNobU4,139190
+onnx_ir/__init__.py,sha256=N-6BRNeRjjU17-iaezMOq1ErVq8irJltkubTpAbZ2BY,3441
+onnx_ir/_core.py,sha256=9AnsXuoiFzypp9U_al7yU46OU0iqrGq0rJzLw4FHmcI,143517
 onnx_ir/_display.py,sha256=230bMN_hVy47Ug3HkA4o5Tf5Hr21AnBEoq5w0fxjyTs,1300
-onnx_ir/_enums.py,sha256=NvfBARtUGc6_qrXRvWdenuhGR7DqDrqigbWaECugTv8,14043
+onnx_ir/_enums.py,sha256=E7WQ7yQzulBeimamc9q_k4fEUoyH_2PWtaOMpwck_W0,13915
 onnx_ir/_graph_comparison.py,sha256=8_D1gu547eCDotEUqxfIJhUGU_Ufhfji7sfsSraOj3g,727
 onnx_ir/_graph_containers.py,sha256=PRKrshRZ5rzWCgRs1TefzJq9n8wyo7OqeKy3XxMhyys,14265
 onnx_ir/_io.py,sha256=GWwA4XOZ-ZX1cgibgaYD0K0O5d9LX21ZwcBN02Wrh04,5205
@@ -16,13 +16,13 @@ onnx_ir/_version_utils.py,sha256=bZThuE7meVHFOY1DLsmss9WshVIp9iig7udGfDbVaK4,133
 onnx_ir/convenience.py,sha256=0B1epuXZCSmY4FbW2vaYfR-t5ubxBZ1UruiytHs-zFw,917
 onnx_ir/external_data.py,sha256=rXHtRU-9tjAt10Iervhr5lsI6Dtv-EhR7J4brxppImA,18079
 onnx_ir/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-onnx_ir/serde.py,sha256=H9V9dff3UF2e-vnsQt9FlUJMnlKWKOms6CjivMtD2c4,78226
+onnx_ir/serde.py,sha256=S0zCZnfePs1UV927HDEr3VnXuue_B5PD1dVPqwuwrak,80636
 onnx_ir/tape.py,sha256=4FyfAHmVhQoMsfHMYnBwP2azi6UF6b6pj--ercObqZs,350
 onnx_ir/tensor_adapters.py,sha256=YffUeZDZi8thxm-4nF2cL6cNSJSVmLm4A3IbEzwY8QQ,7233
 onnx_ir/testing.py,sha256=WTrjf2joWizDWaYMJlV1KjZMQw7YmZ8NvuBTVn1uY6s,8803
 onnx_ir/traversal.py,sha256=Wy4XphwuapAvm94-5iaz6G8LjIoMFpY7qfPfXzYViEE,4488
-onnx_ir/_convenience/__init__.py,sha256=bXUxjZ_91idQJ33zWtByQ0J4VsWCUdvAy9iIflpLtW8,19754
-onnx_ir/_convenience/_constructors.py,sha256=5GhlYy_xCE2ng7l_4cNx06WQsNDyvS-0U1HgOpPKJEk,8347
+onnx_ir/_convenience/__init__.py,sha256=SO7kc8RXVKEUODGh0q2Y7WgmbUsOjYSixmKFx_A0DAQ,19752
+onnx_ir/_convenience/_constructors.py,sha256=HqCGtNPMzFFEerC7I5VEyMdBuIdOJDucn9UEdwuymcg,11519
 onnx_ir/_thirdparty/asciichartpy.py,sha256=afQ0fsqko2uYRPAR4TZBrQxvCb4eN8lxZ2yDFbVQq_s,10533
 onnx_ir/passes/__init__.py,sha256=M_Tcl_-qGSNPluFIvOoeDyh0qAwNayaYyXDS5UJUJPQ,764
 onnx_ir/passes/_pass_infra.py,sha256=xIOw_zZIuOqD4Z_wZ4OvsqXfh2IZMoMlDp1xQ_MPQlc,9567
@@ -30,17 +30,17 @@ onnx_ir/passes/common/__init__.py,sha256=NYBrXvkq_CbWRwcZptSsF4u_-1zfN_BClLhVQY0
 onnx_ir/passes/common/_c_api_utils.py,sha256=g6riA6xNGVWaO5YjVHZ0krrfslWHmRlryRkwB8X56cg,2907
 onnx_ir/passes/common/clear_metadata_and_docstring.py,sha256=YwouLfsNFSaTuGd7uMOGjdvVwG9yHQTkSphUgDlM0ME,2365
 onnx_ir/passes/common/common_subexpression_elimination.py,sha256=wZ1zEPdCshYB_ifP9fCAVfzQkesE6uhCfzCuL2qO5fA,7948
-onnx_ir/passes/common/constant_manipulation.py,sha256=dFzzqbpRecJJrYf6edvR_sdr4F0gV-1wEtDXsQ7fStM,9101
+onnx_ir/passes/common/constant_manipulation.py,sha256=Ja_2uO59ni8YX600iF3wvCmOk4EvM5crl7csYU1s7rQ,9391
 onnx_ir/passes/common/identity_elimination.py,sha256=wN8g8uPGn6IIQ6Jf1lo6nGTXvpWyiSQtT_CfmtvZpwA,3664
 onnx_ir/passes/common/initializer_deduplication.py,sha256=gKrXTMFAtCkMmiIm8zWzwPnwSbRdZxunJeAt_jFU-vY,7253
 onnx_ir/passes/common/inliner.py,sha256=wBoO6yXt6F1AObQjYZHMQ0wn3YH681N4HQQVyaMAYd4,13702
-onnx_ir/passes/common/naming.py,sha256=NNKc9IPrmzm3J0zGQILfooayVzfdXDYHY9DHex1hFgs,10927
+onnx_ir/passes/common/naming.py,sha256=l_LrUiI3gAoSEVs8YeQ5kRNWp7aMOlBK-SfFUsKobZI,10687
 onnx_ir/passes/common/onnx_checker.py,sha256=_sPmJ2ff9pDB1g9q7082BL6fyubomRaj6svE0cCyDew,1691
 onnx_ir/passes/common/shape_inference.py,sha256=LVdvxjeKtcIEbPcb6mKisxoPJOOawzsm3tzk5j9xqeM,3992
 onnx_ir/passes/common/topological_sort.py,sha256=Vcu1YhBdfRX4LROr0NScjB1Pwz2DjBFD0Z_GxqaxPF8,999
 onnx_ir/passes/common/unused_removal.py,sha256=cBNqaqGnUVyCWxsD7hBzYk4qSglVPo3SmHAvkUo5-Oc,7613
-onnx_ir-0.1.8.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-onnx_ir-0.1.8.dist-info/METADATA,sha256=trHcb8JDn1IJRC3B9iLv8o6QjAC7G6etJJw6gfGQE2s,3597
-onnx_ir-0.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnx_ir-0.1.8.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
-onnx_ir-0.1.8.dist-info/RECORD,,
+onnx_ir-0.1.10.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+onnx_ir-0.1.10.dist-info/METADATA,sha256=3ou4NcDwiEwn172tRdD57Tb7y7rCGUyfhFDqan7qM3s,3605
+onnx_ir-0.1.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnx_ir-0.1.10.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
+onnx_ir-0.1.10.dist-info/RECORD,,