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

onnx_ir/serde.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
 """Serialize and deserialize the intermediate representation to/from ONNX protos."""
 
 # NOTES for developers:
@@ -14,12 +14,14 @@
 from __future__ import annotations
 
 import functools
+import typing
 
 __all__ = [
     # Tensors
     "TensorProtoTensor",
     # Deserialization
     "from_proto",
+    "from_onnx_text",
     "deserialize_attribute",
     "deserialize_dimension",
     "deserialize_function",
@@ -29,6 +31,7 @@ __all__ = [
     "deserialize_node",
     "deserialize_opset_import",
     "deserialize_tensor",
+    "deserialize_tensor_shape",
     "deserialize_type_proto_for_shape",
     "deserialize_type_proto_for_type",
     "deserialize_value_info_proto",
@@ -59,14 +62,14 @@ __all__ = [
 import collections
 import logging
 import os
-import typing
-from typing import Any, Callable, List, Mapping, Sequence
+from collections.abc import Mapping, Sequence
+from typing import Any, Callable
 
 import numpy as np
 import onnx
 import onnx.external_data_helper
 
-from onnx_ir import _core, _enums, _metadata, _protocols, _type_casting
+from onnx_ir import _core, _enums, _protocols, _type_casting
 
 if typing.TYPE_CHECKING:
     import google.protobuf.internal.containers as proto_containers
@@ -74,12 +77,11 @@ if typing.TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-_PLEASE_CONTRIBUTE = (
-    "Please contribute by creating a PR at https://github.com/microsoft/onnxscript."
-)
+_PLEASE_CONTRIBUTE = "Please contribute by creating a PR at https://github.com/onnx/onnx-ir."
 _FUNCTION_VALUE_INFO_SUPPORTED_VERSION = (
     10  # ONNX IR version where value info in functions was introduced
 )
+_QUANT_PARAMETER_TENSOR_NAMES_FIELD = "quant_parameter_tensor_names"
 _T = typing.TypeVar("_T", bound=Callable[..., Any])
 
 
@@ -121,16 +123,35 @@ def _unflatten_complex(
     return array[::2] + 1j * array[1::2]
 
 
-def from_proto(
-    proto: onnx.ModelProto
-    | onnx.GraphProto
-    | onnx.NodeProto
-    | onnx.TensorProto
-    | onnx.AttributeProto
-    | onnx.ValueInfoProto
-    | onnx.TypeProto
-    | onnx.FunctionProto,
-) -> Any:
+@typing.overload
+def from_proto(proto: onnx.ModelProto) -> _core.Model: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.GraphProto) -> _core.Graph: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.NodeProto) -> _core.Node: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.TensorProto) -> _protocols.TensorProtocol: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.AttributeProto) -> _core.Attr: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.ValueInfoProto) -> _core.Value: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.TypeProto) -> _core.TypeAndShape: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.FunctionProto) -> _core.Function: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: onnx.TensorShapeProto) -> _core.Shape: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(  # type: ignore[overload-overlap]
+    proto: onnx.TensorShapeProto.Dimension,
+) -> tuple[int | _core.SymbolicDim, str | None]: ...
+@typing.overload
+def from_proto(proto: Sequence[onnx.OperatorSetIdProto]) -> dict[str, int]: ...  # type: ignore[overload-overlap]
+@typing.overload
+def from_proto(proto: Sequence[onnx.StringStringEntryProto]) -> dict[str, str]: ...  # type: ignore[overload-overlap]
+
+
+def from_proto(proto: object) -> object:
     """Deserialize an ONNX proto message to an IR object."""
     if isinstance(proto, onnx.ModelProto):
         return deserialize_model(proto)
@@ -151,24 +172,56 @@ def from_proto(
         )
     if isinstance(proto, onnx.FunctionProto):
         return deserialize_function(proto)
+    if isinstance(proto, onnx.TensorShapeProto):
+        return deserialize_tensor_shape(proto)
+    if isinstance(proto, onnx.TensorShapeProto.Dimension):
+        return deserialize_dimension(proto)
+    if isinstance(proto, Sequence) and all(
+        isinstance(p, onnx.OperatorSetIdProto) for p in proto
+    ):
+        return deserialize_opset_import(proto)
+    if isinstance(proto, Sequence) and all(
+        isinstance(p, onnx.StringStringEntryProto) for p in proto
+    ):
+        return deserialize_metadata_props(proto)
     raise NotImplementedError(
         f"Deserialization of {type(proto)} in from_proto is not implemented. "
         "Use a specific ir.serde.deserialize* function instead."
     )
 
 
-def to_proto(
-    ir_object: _protocols.ModelProtocol
-    | _protocols.GraphProtocol
-    | _protocols.NodeProtocol
-    | _protocols.ValueProtocol
-    | _protocols.AttributeProtocol
-    | _protocols.ReferenceAttributeProtocol
-    | _protocols.TensorProtocol
-    | _protocols.TypeProtocol
-    | _protocols.GraphViewProtocol
-    | _protocols.FunctionProtocol,
-) -> Any:
+def from_onnx_text(model_text: str, /) -> _core.Model:
+    """Convert the ONNX textual representation to an IR model.
+
+    Read more about the textual representation at: https://onnx.ai/onnx/repo-docs/Syntax.html
+    """
+    proto = onnx.parser.parse_model(model_text)
+    return deserialize_model(proto)
+
+
+@typing.overload
+def to_proto(ir_object: _protocols.ModelProtocol) -> onnx.ModelProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.GraphProtocol) -> onnx.GraphProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.NodeProtocol) -> onnx.NodeProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.TensorProtocol) -> onnx.TensorProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.AttributeProtocol) -> onnx.AttributeProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.ReferenceAttributeProtocol) -> onnx.AttributeProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.ValueProtocol) -> onnx.ValueInfoProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.TypeProtocol) -> onnx.TypeProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.FunctionProtocol) -> onnx.FunctionProto: ...  # type: ignore[overload-overlap]
+@typing.overload
+def to_proto(ir_object: _protocols.GraphViewProtocol) -> onnx.GraphProto: ...  # type: ignore[overload-overlap]
+
+
+def to_proto(ir_object: object) -> object:
     """Serialize an IR object to a proto."""
     if isinstance(ir_object, _protocols.ModelProtocol):
         return serialize_model(ir_object)
@@ -180,9 +233,10 @@ def to_proto(
         return serialize_tensor(ir_object)
     if isinstance(ir_object, _protocols.ValueProtocol):
         return serialize_value(ir_object)
-    if isinstance(ir_object, _protocols.AttributeProtocol):
+    if isinstance(ir_object, _protocols.AttributeProtocol) and not ir_object.is_ref():
         return serialize_attribute(ir_object)
     if isinstance(ir_object, _protocols.ReferenceAttributeProtocol):
+        assert ir_object.is_ref()
         return serialize_reference_attribute_into(onnx.AttributeProto(), ir_object)
     if isinstance(ir_object, _protocols.TypeProtocol):
         return serialize_type_into(onnx.TypeProto(), ir_object)
@@ -199,12 +253,11 @@ def to_proto(
 class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
     """A tensor initialized from a tensor proto."""
 
+    __slots__ = ("_proto",)
+
     def __init__(self, proto: onnx.TensorProto) -> None:
+        super().__init__(metadata_props=deserialize_metadata_props(proto.metadata_props))
         self._proto = proto
-        self._metadata_props: dict[str, str] | None = deserialize_metadata_props(
-            proto.metadata_props
-        )
-        self._metadata: _metadata.MetadataStore | None = None
 
     @property
     def name(self) -> str:
@@ -225,7 +278,7 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
     def dtype(self) -> _enums.DataType:
         return _enums.DataType(self._proto.data_type)
 
-    @property
+    @property  # type: ignore[misc]
     def doc_string(self) -> str:
         return self._proto.doc_string
 
@@ -234,9 +287,10 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
         return self._proto
 
     def __repr__(self) -> str:
-        # It is a little hard to display the content when there can be types
-        # unsupported by numpy
-        # Preferably we should display some content when the tensor is small
+        if self.size <= 10:
+            tensor_lines = repr(self.numpy()).split("\n")
+            tensor_text = " ".join(line.strip() for line in tensor_lines)
+            return f"{self._repr_base()}({tensor_text}, name={self.name!r})"
         return f"{self._repr_base()}(name={self.name!r})"
 
     def __array__(self, dtype: Any = None) -> np.ndarray:
@@ -277,8 +331,7 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
             raise ValueError("Cannot convert UNDEFINED tensor to numpy array.")
         if self._proto.data_location == onnx.TensorProto.EXTERNAL:
             raise ValueError(
-                "Cannot convert external tensor to numpy array. "
-                "Use ir.ExternalTensor instead."
+                "Cannot convert external tensor to numpy array. Use ir.ExternalTensor instead."
             )
 
         if self._proto.HasField("raw_data"):
@@ -323,6 +376,8 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
             return _type_casting.unpack_int4(array.astype(np.uint8), self._proto.dims)
         elif dtype == _enums.DataType.UINT4:
             return _type_casting.unpack_uint4(array.astype(np.uint8), self._proto.dims)
+        elif dtype == _enums.DataType.FLOAT4E2M1:
+            return _type_casting.unpack_float4e2m1(array.astype(np.uint8), self._proto.dims)
         else:
             # Otherwise convert to the correct dtype and reshape
             # Note we cannot use view() here because the storage dtype may not be the same size as the target
@@ -369,6 +424,7 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
                 _enums.DataType.FLOAT8E5M2FNUZ,
                 _enums.DataType.INT4,
                 _enums.DataType.UINT4,
+                _enums.DataType.FLOAT4E2M1,
             }:
                 # uint4 and int4 values are already packed, even when stored as int32
                 # so we don't need to pack them again
@@ -393,23 +449,6 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
         # For example, int32_data can be empty and still be a valid tensor.
         return b""
 
-    @property
-    def meta(self) -> _metadata.MetadataStore:
-        """The metadata store for intermediate analysis.
-
-        Write to the :attr:`metadata_props` if you would like the metadata to be serialized
-        to the ONNX proto.
-        """
-        if self._metadata is None:
-            self._metadata = _metadata.MetadataStore()
-        return self._metadata
-
-    @property
-    def metadata_props(self) -> dict[str, str]:
-        if self._metadata_props is None:
-            self._metadata_props = {}
-        return self._metadata_props
-
 
 def _get_field(proto: Any, field: str) -> Any:
     if proto.HasField(field):
@@ -472,7 +511,7 @@ def deserialize_model(proto: onnx.ModelProto) -> _core.Model:
         model_version=_get_field(proto, "model_version"),
         doc_string=_get_field(proto, "doc_string"),
         functions=functions,
-        meta_data_props=deserialize_metadata_props(proto.metadata_props),
+        metadata_props=deserialize_metadata_props(proto.metadata_props),
     )
 
     # Handle experimental value info for functions created by the dynamo exporter in IR version 9
@@ -541,6 +580,9 @@ def deserialize_graph(proto: onnx.GraphProto) -> _core.Graph:
 
     Returns:
         IR Graph.
+
+    .. versionadded:: 0.3
+        Support for *quantization_annotation* is added.
     """
     return _deserialize_graph(proto, [])
 
@@ -561,44 +603,89 @@ def _deserialize_graph(
     Returns:
         IR Graph.
     """
+    # Process TensorAnnotation for quantization
+    quantization_annotations = {
+        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]
     for info, value in zip(proto.input, inputs):
         deserialize_value_info_proto(info, value)
 
+        # Add TensorAnnotation for inputs if they exist
+        if value.name in quantization_annotations:
+            _deserialize_quantization_annotation(quantization_annotations[value.name], value)
+
     # Initialize the values dictionary for this graph scope with the inputs and initializers
     values: dict[str, _core.Value] = {v.name: v for v in inputs}  # type: ignore[misc]
+
+    # Enter the graph scope by pushing the values for this scope to the stack
     scoped_values.append(values)
+
     initializer_values = []
-    for tensor in initializer_tensors:
-        if tensor.name in values:
+    for i, tensor in enumerate(initializer_tensors):
+        initializer_name = tensor.name
+        if not initializer_name:
+            logger.warning(
+                "Initializer tensor must have a name but the %s-th initializer does not. Skipping this initializer.",
+                i,
+            )
+            continue
+        if initializer_name in values:
             # The initializer is for an input
-            initializer_value = values[tensor.name]
+            initializer_value = values[initializer_name]
             initializer_value.const_value = tensor
         else:
             # The initializer is for some other value. Create this value first
             initializer_value = _core.Value(
                 None,
                 index=None,
-                name=tensor.name,
-                # TODO(justinchuby): Fix type hinting for shape and dtype
-                shape=tensor.shape,  # type: ignore
+                name=initializer_name,
+                # Include shape and type even if the shape or type is not provided as ValueInfoProto.
+                # Users expect initialized values to have shape and type information.
                 type=_core.TensorType(tensor.dtype),
+                shape=tensor.shape,  # type: ignore[arg-type]
                 const_value=tensor,
             )
-            values[tensor.name] = initializer_value  # type: ignore[index]
+            if initializer_value.name in quantization_annotations:
+                _deserialize_quantization_annotation(
+                    quantization_annotations[initializer_value.name], initializer_value
+                )
+            values[initializer_name] = initializer_value
         initializer_values.append(initializer_value)
 
-    # Add ValueInfos for this graph scope
+    # 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}
 
     # Deserialize nodes with all known values
-    nodes = [_deserialize_node(node, scoped_values, value_info) for node in proto.node]
+    nodes = [
+        _deserialize_node(node, scoped_values, value_info, quantization_annotations)
+        for node in proto.node
+    ]
 
-    # Fill in values for graph outputs
-    outputs = [deserialize_value_info_proto(info, values[info.name]) for info in proto.output]
+    outputs = []
+    for info in proto.output:
+        # Fill in values for graph outputs
+        output_name = info.name
+        if output_name not in values:
+            # Handle (invalid) graph outputs that do not have any producers
+            logger.warning(
+                "Output '%s' is not produced by any node. The graph has an invalid output",
+                output_name,
+            )
+            value = _core.Value(name=output_name)
+        else:
+            # A valid, normal graph output
+            value = values[output_name]
+        # Fill in shape/type information
+        deserialize_value_info_proto(info, value)
+        outputs.append(value)
+
+    # Exit the graph scope by popping the values for this scope from the stack
     scoped_values.pop()
+
     return _core.Graph(
         inputs,
         outputs,
@@ -617,7 +704,10 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
     value_info = {info.name: info for info in getattr(proto, "value_info", [])}
 
     # TODO(justinchuby): Handle unsorted nodes
-    nodes = [_deserialize_node(node, [values], value_info=value_info) for node in proto.node]
+    nodes = [
+        _deserialize_node(node, [values], value_info=value_info, quantization_annotations={})
+        for node in proto.node
+    ]
     outputs = [values[name] for name in proto.output]
     graph = _core.Graph(
         inputs,
@@ -631,6 +721,7 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
             if hasattr(proto, "overload") and proto.overload
             else ""
         ),
+        metadata_props=deserialize_metadata_props(proto.metadata_props),
     )
     attributes = [_deserialize_attribute(attr, []) for attr in proto.attribute_proto]
     # Attributes without defaults
@@ -642,8 +733,7 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
         name=proto.name,
         overload=getattr(proto, "overload", ""),
         graph=graph,
-        attributes=typing.cast(List[_core.Attr], attributes),
-        metadata_props=deserialize_metadata_props(proto.metadata_props),
+        attributes=attributes,
     )
 
 
@@ -662,29 +752,41 @@ def deserialize_value_info_proto(
     return value
 
 
+@_capture_errors(lambda proto, value: str(proto))
+def _deserialize_quantization_annotation(
+    proto: onnx.TensorAnnotation, value: _core.Value
+) -> None:
+    """Deserialize a quantization_annotation as TensorAnnotation into a Value.
+
+    This function is marked private because we don't expect users to call it directly.
+    """
+    value.meta[_QUANT_PARAMETER_TENSOR_NAMES_FIELD] = _deserialize_string_string_maps(
+        proto.quant_parameter_tensor_names
+    )
+
+
+@_capture_errors(str)
+def deserialize_tensor_shape(proto: onnx.TensorShapeProto) -> _core.Shape:
+    # This logic handles when the shape is [] as well
+    dim_protos = proto.dim
+    deserialized_dim_denotations = [
+        deserialize_dimension(dim_proto) for dim_proto in dim_protos
+    ]
+    dims = [dim for dim, _ in deserialized_dim_denotations]
+    denotations = [denotation for _, denotation in deserialized_dim_denotations]
+    return _core.Shape(dims, denotations=denotations, frozen=True)
+
+
 @_capture_errors(str)
 def deserialize_type_proto_for_shape(proto: onnx.TypeProto) -> _core.Shape | None:
     if proto.HasField("tensor_type"):
         if (shape_proto := _get_field(proto.tensor_type, "shape")) is None:
             return None
-        # This logic handles when the shape is [] as well
-        dim_protos = shape_proto.dim
-        deserialized_dim_denotations = [
-            deserialize_dimension(dim_proto) for dim_proto in dim_protos
-        ]
-        dims = [dim for dim, _ in deserialized_dim_denotations]
-        denotations = [denotation for _, denotation in deserialized_dim_denotations]
-        return _core.Shape(dims, denotations=denotations, frozen=True)
+        return deserialize_tensor_shape(shape_proto)
     if proto.HasField("sparse_tensor_type"):
         if (shape_proto := _get_field(proto.sparse_tensor_type, "shape")) is None:
             return None
-        dim_protos = shape_proto.dim
-        deserialized_dim_denotations = [
-            deserialize_dimension(dim_proto) for dim_proto in dim_protos
-        ]
-        dims = [dim for dim, _ in deserialized_dim_denotations]
-        denotations = [denotation for _, denotation in deserialized_dim_denotations]
-        return _core.Shape(dims, denotations=denotations, frozen=True)
+        return deserialize_tensor_shape(shape_proto)
     if proto.HasField("sequence_type"):
         if (elem_type := _get_field(proto.sequence_type, "elem_type")) is None:
             return None
@@ -800,14 +902,17 @@ def deserialize_metadata_props(
     return {entry.key: entry.value for entry in proto}
 
 
-def deserialize_attribute(proto: onnx.AttributeProto) -> _core.Attr | _core.RefAttr:
+_deserialize_string_string_maps = deserialize_metadata_props
+
+
+def deserialize_attribute(proto: onnx.AttributeProto) -> _core.Attr:
     return _deserialize_attribute(proto, [])
 
 
 @_capture_errors(lambda proto, scoped_values: str(proto))
 def _deserialize_attribute(
     proto: onnx.AttributeProto, scoped_values: list[dict[str, _core.Value]]
-) -> _core.Attr | _core.RefAttr:
+) -> _core.Attr:
     name = proto.name
     doc_string = _get_field(proto, "doc_string")
     type_ = _enums.AttributeType(proto.type)
@@ -874,14 +979,17 @@ def _deserialize_attribute(
 
 
 def deserialize_node(proto: onnx.NodeProto) -> _core.Node:
-    return _deserialize_node(proto, scoped_values=[], value_info={})
+    return _deserialize_node(
+        proto, scoped_values=[{}], value_info={}, quantization_annotations={}
+    )
 
 
-@_capture_errors(lambda proto, scoped_values, value_info: str(proto))
+@_capture_errors(lambda proto, scoped_values, value_info, quantization_annotations: str(proto))
 def _deserialize_node(
     proto: onnx.NodeProto,
     scoped_values: list[dict[str, _core.Value]],
     value_info: dict[str, onnx.ValueInfoProto],
+    quantization_annotations: dict[str, onnx.TensorAnnotation],
 ) -> _core.Node:
     node_inputs: list[_core.Value | None] = []
     for input_name in proto.input:
@@ -924,6 +1032,10 @@ def _deserialize_node(
             # Fill in shape/type information if they exist
             if input_name in value_info:
                 deserialize_value_info_proto(value_info[input_name], value)
+            if input_name in quantization_annotations:
+                _deserialize_quantization_annotation(
+                    quantization_annotations[input_name], value
+                )
             node_inputs.append(value)
             # We can only create the value in the current scope. If the subgraph is
             # referencing a value that is not in the current scope, it is impossible
@@ -965,6 +1077,8 @@ def _deserialize_node(
                 proto.name,
                 proto.op_type,
             )
+        if output_name in quantization_annotations:
+            _deserialize_quantization_annotation(quantization_annotations[output_name], value)
         node_outputs.append(value)
     return _core.Node(
         proto.domain,
@@ -1036,7 +1150,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
-    return not (value.shape is None and value.type is None)
+    if value.shape is None and value.type is None:
+        return False
+    if not value.name:
+        logger.debug("Did not serialize '%s' because its name is empty", value)
+        return False
+    return True
 
 
 def _serialize_experimental_value_info_for_function_ir9_into(
@@ -1063,7 +1182,7 @@ def _serialize_experimental_value_info_for_function_ir9_into(
 
     for input in function.inputs:
         if not input.name:
-            logging.warning(
+            logger.warning(
                 "Function '%s': Value name not set for function input: %s",
                 function_qualified_name,
                 input,
@@ -1076,7 +1195,7 @@ def _serialize_experimental_value_info_for_function_ir9_into(
     for node in function:
         for node_output in node.outputs:
             if not node_output.name:
-                logging.warning(
+                logger.warning(
                     "Function '%s': Value name not set for node output: %s",
                     function_qualified_name,
                     node_output,
@@ -1107,23 +1226,46 @@ def _serialize_opset_imports_into(
         opset_ids.add(domain=domain, version=version)
 
 
-def _serialize_metadata_props_into(
+def _serialize_string_string_maps(
     string_string_entries: proto_containers.RepeatedCompositeFieldContainer[
         onnx.StringStringEntryProto
     ],
     from_: Mapping[str, str],
 ) -> None:
-    """Serialize metadata properties into a repeated field of string-string entries.
+    """Serialize a <str, str> mapping into a repeated field of string-string entries.
 
     Args:
         string_string_entries: The repeated field to serialize into.
-        from_: The mapping of metadata properties to serialize.
+        from_: The mapping of a <str, str> mapping to serialize.
     """
     # Sort names for deterministic serialization
     for key in sorted(from_):
         string_string_entries.add(key=key, value=from_[key])
 
 
+_serialize_metadata_props_into = _serialize_string_string_maps
+
+
+def _maybe_add_quantization_annotation(
+    graph_proto: onnx.GraphProto, value: _protocols.ValueProtocol
+) -> None:
+    if quantization_annotation := value.meta.get(_QUANT_PARAMETER_TENSOR_NAMES_FIELD):
+        _serialize_tensor_annotation_into(
+            graph_proto.quantization_annotation.add(), value.name, quantization_annotation
+        )
+
+
+def _serialize_tensor_annotation_into(
+    tensor_annotation_proto: onnx.TensorAnnotation,
+    tensor_name: str,
+    quant_parameter_tensor_names: dict[str, str],
+) -> None:
+    tensor_annotation_proto.tensor_name = tensor_name
+    _serialize_string_string_maps(
+        tensor_annotation_proto.quant_parameter_tensor_names, quant_parameter_tensor_names
+    )
+
+
 def serialize_graph(
     graph: _protocols.GraphProtocol | _protocols.GraphViewProtocol,
 ) -> onnx.GraphProto:
@@ -1159,29 +1301,41 @@ def serialize_graph_into(
         graph_proto.doc_string = from_.doc_string
     for input_ in from_.inputs:
         serialize_value_into(graph_proto.input.add(), input_)
+        if input_.name not in from_.initializers:
+            # Annotations for initializers will be added below to avoid double adding
+            # TODO(justinchuby): We should add a method is_initializer() on Value when
+            # the initializer list is tracked
+            _maybe_add_quantization_annotation(graph_proto, input_)
+    input_names = {input_.name for input_ in from_.inputs}
     # TODO(justinchuby): Support sparse_initializer
-    for initializer in from_.initializers.values():
-        if initializer.const_value is None:
+    for value in from_.initializers.values():
+        _maybe_add_quantization_annotation(graph_proto, value)
+        if _should_create_value_info_for_value(value) and value.name not in input_names:
+            # Serialize information about all initializers into value_info,
+            # except for those that are also graph inputs
+            serialize_value_into(graph_proto.value_info.add(), value)
+        if value.const_value is None:
             # Skip initializers without constant values
-            logger.warning(
-                "Initializer '%s' does not have a constant value set.", initializer.name
-            )
+            logger.warning("Initializer '%s' does not have a constant value set.", value.name)
             continue
         # Make sure the tensor's name is the same as the value's name
-        initializer.const_value.name = initializer.name
-        serialize_tensor_into(graph_proto.initializer.add(), from_=initializer.const_value)
+        value.const_value.name = value.name
+        serialize_tensor_into(graph_proto.initializer.add(), from_=value.const_value)
     for node in from_:
         serialize_node_into(graph_proto.node.add(), from_=node)
         for node_output in node.outputs:
-            if not _should_create_value_info_for_value(node_output):
-                # No need to serialize value info if it is not set
-                continue
             if node_output.is_graph_output():
-                # No need to serialize value info for these outputs because they are also graph outputs
+                # No need to serialize info for these outputs because they are handled as graph outputs
                 continue
-            serialize_value_into(graph_proto.value_info.add(), node_output)
+            _maybe_add_quantization_annotation(graph_proto, node_output)
+            if not _should_create_value_info_for_value(node_output):  # pylint: disable=no-else-continue
+                # No need to serialize value info if it is not set
+                continue
+            else:
+                serialize_value_into(graph_proto.value_info.add(), node_output)
     for output in from_.outputs:
         serialize_value_into(graph_proto.output.add(), from_=output)
+        _maybe_add_quantization_annotation(graph_proto, output)
     if from_.metadata_props:
         _serialize_metadata_props_into(graph_proto.metadata_props, from_.metadata_props)
 
@@ -1269,6 +1423,23 @@ def serialize_node(node: _protocols.NodeProtocol) -> onnx.NodeProto:
     return node_proto
 
 
+def _remove_trailing_outputs(
+    outputs: Sequence[_protocols.ValueProtocol],
+) -> Sequence[_protocols.ValueProtocol]:
+    """Remove trailing outputs that have empty names.
+
+    Args:
+        outputs: The outputs to remove trailing outputs from.
+
+    Returns:
+        The outputs with trailing outputs removed.
+    """
+    for i, output in enumerate(reversed(outputs)):
+        if output.name:
+            return outputs[: len(outputs) - i]
+    return []
+
+
 @_capture_errors(lambda node_proto, from_: repr(from_))
 def serialize_node_into(node_proto: onnx.NodeProto, from_: _protocols.NodeProtocol) -> None:
     node_proto.op_type = from_.op_type
@@ -1288,23 +1459,16 @@ def serialize_node_into(node_proto: onnx.NodeProto, from_: _protocols.NodeProtoc
             node_proto.input.append("")
         else:
             node_proto.input.append(input_.name)
-    for output in from_.outputs:
+
+    # Do not include the trailing outputs that have empty names
+    for output in _remove_trailing_outputs(from_.outputs):
         node_proto.output.append(output.name)
+
     for attr in from_.attributes.values():
-        if isinstance(attr, _core.Attr):
-            serialize_attribute_into(node_proto.attribute.add(), from_=attr)
-        elif isinstance(attr, _core.RefAttr):
-            serialize_reference_attribute_into(node_proto.attribute.add(), from_=attr)
-        # Handle protocol attributes for completeness. We do not check them first because
-        # calling isinstance on a protocol can be slow.
-        # Most of the time, we will have Attr or RefAttr so the two branches below
-        # will not be taken.
-        elif isinstance(attr, _protocols.AttributeProtocol):
-            serialize_attribute_into(node_proto.attribute.add(), from_=attr)
-        elif isinstance(attr, _protocols.ReferenceAttributeProtocol):
-            serialize_reference_attribute_into(node_proto.attribute.add(), from_=attr)
+        if not attr.is_ref():
+            serialize_attribute_into(node_proto.attribute.add(), from_=attr)  # type: ignore[arg-type]
         else:
-            raise TypeError(f"Unsupported attribute type: {type(attr)}")
+            serialize_reference_attribute_into(node_proto.attribute.add(), from_=attr)  # type: ignore[arg-type]
 
 
 def serialize_tensor(tensor: _protocols.TensorProtocol) -> onnx.TensorProto: