onnx 1.16.1__cp311-cp311-win32.whl → 1.17.0__cp311-cp311-win32.whl

onnx/external_data_helper.py

@@ -1,15 +1,23 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
 import os
 import re
 import sys
 import uuid
 from itertools import chain
-from typing import Callable, Iterable, Optional
+from typing import Callable, Iterable
 
 import onnx.onnx_cpp2py_export.checker as c_checker
-from onnx.onnx_pb import AttributeProto, GraphProto, ModelProto, TensorProto
+from onnx.onnx_pb import (
+    AttributeProto,
+    FunctionProto,
+    GraphProto,
+    ModelProto,
+    TensorProto,
+)
 
 
 class ExternalDataInfo:
@@ -71,10 +79,10 @@ def load_external_data_for_model(model: ModelProto, base_dir: str) -> None:
 def set_external_data(
     tensor: TensorProto,
     location: str,
-    offset: Optional[int] = None,
-    length: Optional[int] = None,
-    checksum: Optional[str] = None,
-    basepath: Optional[str] = None,
+    offset: int | None = None,
+    length: int | None = None,
+    checksum: str | None = None,
+    basepath: str | None = None,
 ) -> None:
     if not tensor.HasField("raw_data"):
         raise ValueError(
@@ -101,7 +109,7 @@ def set_external_data(
 def convert_model_to_external_data(
     model: ModelProto,
     all_tensors_to_one_file: bool = True,
-    location: Optional[str] = None,
+    location: str | None = None,
     size_threshold: int = 1024,
     convert_attribute: bool = False,
 ) -> None:
@@ -220,11 +228,12 @@ def _recursive_attribute_processor(
 
 
 def _get_initializer_tensors_from_graph(
-    onnx_model_proto_graph: GraphProto,
+    graph_or_function: GraphProto | FunctionProto, /
 ) -> Iterable[TensorProto]:
-    """Create an iterator of initializer tensors from ONNX model graph."""
-    yield from onnx_model_proto_graph.initializer
-    for node in onnx_model_proto_graph.node:
+    """Create an iterator of initializer tensors from ONNX model graph/function."""
+    if isinstance(graph_or_function, GraphProto):
+        yield from graph_or_function.initializer
+    for node in graph_or_function.node:
         for attribute in node.attribute:
             yield from _recursive_attribute_processor(
                 attribute, _get_initializer_tensors_from_graph
@@ -234,13 +243,15 @@ def _get_initializer_tensors_from_graph(
 def _get_initializer_tensors(onnx_model_proto: ModelProto) -> Iterable[TensorProto]:
     """Create an iterator of initializer tensors from ONNX model."""
     yield from _get_initializer_tensors_from_graph(onnx_model_proto.graph)
+    for function in onnx_model_proto.functions:
+        yield from _get_attribute_tensors_from_graph(function)
 
 
 def _get_attribute_tensors_from_graph(
-    onnx_model_proto_graph: GraphProto,
+    graph_or_function: GraphProto | FunctionProto, /
 ) -> Iterable[TensorProto]:
-    """Create an iterator of tensors from node attributes of an ONNX model graph."""
-    for node in onnx_model_proto_graph.node:
+    """Create an iterator of tensors from node attributes of an ONNX model graph/function."""
+    for node in graph_or_function.node:
         for attribute in node.attribute:
             if attribute.HasField("t"):
                 yield attribute.t
@@ -253,6 +264,8 @@ def _get_attribute_tensors_from_graph(
 def _get_attribute_tensors(onnx_model_proto: ModelProto) -> Iterable[TensorProto]:
     """Create an iterator of tensors from node attributes of an ONNX model."""
     yield from _get_attribute_tensors_from_graph(onnx_model_proto.graph)
+    for function in onnx_model_proto.functions:
+        yield from _get_attribute_tensors_from_graph(function)
 
 
 def _is_valid_filename(filename: str) -> bool:

onnx/gen_proto.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import argparse
 import glob
@@ -8,7 +9,7 @@ import os
 import re
 import subprocess
 from textwrap import dedent
-from typing import Iterable, Optional
+from typing import Iterable
 
 autogen_header = """\
 //
@@ -121,7 +122,7 @@ def translate(source: str, proto: int, onnx_ml: bool, package_name: str) -> str:
     return "\n".join(lines)  # TODO: not Windows friendly
 
 
-def qualify(f: str, pardir: Optional[str] = None) -> str:
+def qualify(f: str, pardir: str | None = None) -> str:
     if pardir is None:
         pardir = os.path.realpath(os.path.dirname(__file__))
     return os.path.join(pardir, f)

onnx/helper.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import collections.abc
 import numbers
@@ -12,7 +13,6 @@ from typing import (
     Dict,
     KeysView,
     List,
-    Optional,
     Sequence,
     Tuple,
     TypeVar,
@@ -23,6 +23,7 @@ from typing import (
 import google.protobuf.message
 import numpy as np
 
+import onnx._custom_element_types as custom_np_types
 from onnx import (
     IR_VERSION,
     AttributeProto,
@@ -76,6 +77,7 @@ VERSION_TABLE: VersionTableType = [
     ("1.14.1", 9, 19, 3, 1),
     ("1.15.0", 9, 20, 4, 1),
     ("1.16.0", 10, 21, 5, 1),
+    ("1.17.0", 10, 22, 5, 1),
 ]
 
 VersionMapType = Dict[Tuple[str, int], int]
@@ -116,7 +118,7 @@ def find_min_ir_version_for(
     """
     default_min_version = 3
 
-    def find_min(domain: Union[str, None], version: int) -> int:
+    def find_min(domain: str | None, version: int) -> int:
         key = (domain or "ai.onnx", version)
         if key in OP_SET_ID_VERSION_MAP:
             return OP_SET_ID_VERSION_MAP[key]
@@ -133,10 +135,10 @@ def make_node(
     op_type: str,
     inputs: Sequence[str],
     outputs: Sequence[str],
-    name: Optional[str] = None,
-    doc_string: Optional[str] = None,
-    domain: Optional[str] = None,
-    overload: Optional[str] = None,
+    name: str | None = None,
+    doc_string: str | None = None,
+    domain: str | None = None,
+    overload: str | None = None,
     **kwargs: Any,
 ) -> NodeProto:
     """Construct a NodeProto.
@@ -201,10 +203,10 @@ def make_graph(
     name: str,
     inputs: Sequence[ValueInfoProto],
     outputs: Sequence[ValueInfoProto],
-    initializer: Optional[Sequence[TensorProto]] = None,
-    doc_string: Optional[str] = None,
-    value_info: Optional[Sequence[ValueInfoProto]] = None,
-    sparse_initializer: Optional[Sequence[SparseTensorProto]] = None,
+    initializer: Sequence[TensorProto] | None = None,
+    doc_string: str | None = None,
+    value_info: Sequence[ValueInfoProto] | None = None,
+    sparse_initializer: Sequence[SparseTensorProto] | None = None,
 ) -> GraphProto:
     """Construct a GraphProto
 
@@ -261,11 +263,11 @@ def make_function(
     outputs: Sequence[str],
     nodes: Sequence[NodeProto],
     opset_imports: Sequence[OperatorSetIdProto],
-    attributes: Optional[Sequence[str]] = None,
-    attribute_protos: Optional[Sequence[AttributeProto]] = None,
-    doc_string: Optional[str] = None,
-    overload: Optional[str] = None,
-    value_info: Optional[Sequence[ValueInfoProto]] = None,
+    attributes: Sequence[str] | None = None,
+    attribute_protos: Sequence[AttributeProto] | None = None,
+    doc_string: str | None = None,
+    overload: str | None = None,
+    value_info: Sequence[ValueInfoProto] | None = None,
 ) -> FunctionProto:
     if attributes is None:
         attributes = []
@@ -305,7 +307,7 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
     model.ir_version = IR_VERSION
     model.graph.CopyFrom(graph)
 
-    opset_imports: Optional[Sequence[OperatorSetIdProto]] = None
+    opset_imports: Sequence[OperatorSetIdProto] | None = None
     opset_imports = kwargs.pop("opset_imports", None)  # type: ignore
     if opset_imports is not None:
         model.opset_import.extend(opset_imports)
@@ -314,7 +316,7 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
         imp = model.opset_import.add()
         imp.version = defs.onnx_opset_version()
 
-    functions: Optional[Sequence[FunctionProto]] = None
+    functions: Sequence[FunctionProto] | None = None
     functions = kwargs.pop("functions", None)  # type: ignore
     if functions is not None:
         model.functions.extend(functions)
@@ -337,10 +339,15 @@ def make_model_gen_version(graph: GraphProto, **kwargs: Any) -> ModelProto:
 
 
 def set_metadata_props(
-    proto: Union[
-        ModelProto, GraphProto, FunctionProto, NodeProto, TensorProto, ValueInfoProto
-    ],
-    dict_value: Dict[str, str],
+    proto: (
+        ModelProto
+        | GraphProto
+        | FunctionProto
+        | NodeProto
+        | TensorProto
+        | ValueInfoProto
+    ),
+    dict_value: dict[str, str],
 ) -> None:
     del proto.metadata_props[:]
     for k, v in dict_value.items():
@@ -349,7 +356,7 @@ def set_metadata_props(
         entry.value = v
 
 
-def set_model_props(model: ModelProto, dict_value: Dict[str, str]) -> None:
+def set_model_props(model: ModelProto, dict_value: dict[str, str]) -> None:
     set_metadata_props(model, dict_value)
 
 
@@ -639,9 +646,7 @@ def float32_to_float8e5m2(  # noqa: PLR0911
         raise NotImplementedError("fn and uz must be both False or True.")
 
 
-def pack_float32_to_4bit(
-    array: Union[np.ndarray, Sequence], signed: bool
-) -> np.ndarray:
+def pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndarray:
     """Convert an array of float32 value to a 4bit data-type and pack every two concecutive elements in a byte.
     See :ref:`onnx-detail-int4` for technical details.
 
@@ -770,12 +775,13 @@ def make_tensor(
             TensorProto.INT4,
         ):
             signed = data_type == TensorProto.INT4
-            vals = (
-                pack_float32_to_4bit(vals, signed=signed)
-                .astype(np_dtype)
-                .flatten()
-                .tolist()
-            )
+
+            # Two packed 4-bit values must be represented as a single uint8 value.
+            # Therefore, pack_float32_to_4bit() sets the dtype of the output vals
+            # to uint8 regardless of the value of 'signed'. Using int8 would cause
+            # the size of int4 tensors to increase ~5x if the tensor contains negative values (due to
+            # the way negative values are serialized by protobuf).
+            vals = pack_float32_to_4bit(vals, signed=signed).flatten().tolist()
         elif data_type == TensorProto.BOOL:
             vals = np.array(vals).astype(int)
         elif data_type == TensorProto.STRING:
@@ -836,7 +842,7 @@ def make_sequence(
 
 
 def make_map(
-    name: str, key_type: int, keys: List[Any], values: SequenceProto
+    name: str, key_type: int, keys: list[Any], values: SequenceProto
 ) -> MapProto:
     """Make a Map with specified key-value pair arguments.
 
@@ -869,7 +875,7 @@ def make_map(
 def make_optional(
     name: str,
     elem_type: OptionalProto.DataType,
-    value: Optional[Any],
+    value: Any | None,
 ) -> OptionalProto:
     """Make an Optional with specified value arguments."""
     optional = OptionalProto()
@@ -895,7 +901,7 @@ def make_optional(
     return optional
 
 
-def _to_bytes(value: Union[str, bytes]) -> bytes:
+def _to_bytes(value: str | bytes) -> bytes:
     """Coerce a string (or bytes) value into UTF-8 bytes."""
     return value if isinstance(value, bytes) else value.encode("utf-8")
 
@@ -903,8 +909,8 @@ def _to_bytes(value: Union[str, bytes]) -> bytes:
 def make_attribute(
     key: str,
     value: Any,
-    doc_string: Optional[str] = None,
-    attr_type: Optional[int] = None,
+    doc_string: str | None = None,
+    attr_type: int | None = None,
 ) -> AttributeProto:
     """Makes an AttributeProto based on the value type."""
     attr = AttributeProto()
@@ -995,7 +1001,7 @@ def make_attribute(
 
 
 def make_attribute_ref(
-    name: str, attr_type: AttributeProto.AttributeType, doc_string: Optional[str] = None
+    name: str, attr_type: AttributeProto.AttributeType, doc_string: str | None = None
 ) -> AttributeProto:
     """Make an AttributeProto holding a reference to the parent function's attribute of given name and type."""
     attr = AttributeProto()
@@ -1059,8 +1065,8 @@ def make_empty_tensor_value_info(name: str) -> ValueInfoProto:
 
 def make_tensor_type_proto(
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
-    shape_denotation: Optional[List[str]] = None,
+    shape: Sequence[str | int | None] | None,
+    shape_denotation: list[str] | None = None,
 ) -> TypeProto:
     """Makes a Tensor TypeProto based on the data type and shape."""
     type_proto = TypeProto()
@@ -1105,9 +1111,9 @@ def make_tensor_type_proto(
 def make_tensor_value_info(
     name: str,
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
+    shape: Sequence[str | int | None] | None,
     doc_string: str = "",
-    shape_denotation: Optional[List[str]] = None,
+    shape_denotation: list[str] | None = None,
 ) -> ValueInfoProto:
     """Makes a ValueInfoProto based on the data type and shape."""
     value_info_proto = ValueInfoProto()
@@ -1122,8 +1128,8 @@ def make_tensor_value_info(
 
 def make_sparse_tensor_type_proto(
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
-    shape_denotation: Optional[List[str]] = None,
+    shape: Sequence[str | int | None] | None,
+    shape_denotation: list[str] | None = None,
 ) -> TypeProto:
     """Makes a SparseTensor TypeProto based on the data type and shape."""
     type_proto = TypeProto()
@@ -1168,9 +1174,9 @@ def make_sparse_tensor_type_proto(
 def make_sparse_tensor_value_info(
     name: str,
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
+    shape: Sequence[str | int | None] | None,
     doc_string: str = "",
-    shape_denotation: Optional[List[str]] = None,
+    shape_denotation: list[str] | None = None,
 ) -> ValueInfoProto:
     """Makes a SparseTensor ValueInfoProto based on the data type and shape."""
     value_info_proto = ValueInfoProto()
@@ -1231,7 +1237,7 @@ def make_value_info(
     return value_info_proto
 
 
-def _sanitize_str(s: Union[str, bytes]) -> str:
+def _sanitize_str(s: str | bytes) -> str:
     if isinstance(s, str):
         sanitized = s
     elif isinstance(s, bytes):
@@ -1246,9 +1252,9 @@ def _sanitize_str(s: Union[str, bytes]) -> str:
 def make_tensor_sequence_value_info(
     name: str,
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
+    shape: Sequence[str | int | None] | None,
     doc_string: str = "",
-    elem_shape_denotation: Optional[List[str]] = None,
+    elem_shape_denotation: list[str] | None = None,
 ) -> ValueInfoProto:
     """Makes a Sequence[Tensors] ValueInfoProto based on the data type and shape."""
     value_info_proto = ValueInfoProto()
@@ -1265,7 +1271,7 @@ def make_tensor_sequence_value_info(
 
 def printable_attribute(
     attr: AttributeProto, subgraphs: bool = False
-) -> Union[str, Tuple[str, List[GraphProto]]]:
+) -> str | tuple[str, list[GraphProto]]:
     content = []
     content.append(attr.name)
     content.append("=")
@@ -1381,13 +1387,13 @@ def printable_tensor_proto(t: TensorProto) -> str:
 
 def printable_node(
     node: NodeProto, prefix: str = "", subgraphs: bool = False
-) -> Union[str, Tuple[str, List[GraphProto]]]:
+) -> str | tuple[str, list[GraphProto]]:
     content = []
     if len(node.output):
         content.append(", ".join([f"%{name}" for name in node.output]))
         content.append("=")
     # To deal with nested graphs
-    graphs: List[GraphProto] = []
+    graphs: list[GraphProto] = []
     printed_attrs = []
     for attr in node.attribute:
         if subgraphs:
@@ -1444,7 +1450,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
             content.append(prefix + " ".join(header))
             header = []
             for line in in_strs:
-                content.append(prefix + "  " + line)
+                content.append(prefix + "  " + line)  # noqa: PERF401
         header.append(")")
 
         if in_with_init_strs:
@@ -1452,7 +1458,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
             content.append(prefix + " ".join(header))
             header = []
             for line in in_with_init_strs:
-                content.append(prefix + "  " + line)
+                content.append(prefix + "  " + line)  # noqa: PERF401
             header.append(")")
 
         # from IR 4 onwards an initializer is not required to have a matching graph input
@@ -1468,12 +1474,12 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
             content.append(prefix + " ".join(header))
             header = []
             for line in init_strs:
-                content.append(prefix + "  " + line)
+                content.append(prefix + "  " + line)  # noqa: PERF401
             header.append(")")
 
     header.append("{")
     content.append(prefix + " ".join(header))
-    graphs: List[GraphProto] = []
+    graphs: list[GraphProto] = []
     # body
     for node in graph.node:
         contents_subgraphs = printable_node(node, indent, subgraphs=True)
@@ -1489,7 +1495,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
     # closing bracket
     content.append(prefix + "}")
     for g in graphs:
-        content.append("\n" + printable_graph(g))
+        content.append("\n" + printable_graph(g))  # noqa: PERF401
     return "\n".join(content)
 
 
@@ -1513,8 +1519,8 @@ def strip_doc_string(proto: google.protobuf.message.Message) -> None:
 def make_training_info(
     algorithm: GraphProto,
     algorithm_bindings: AssignmentBindingType,
-    initialization: Optional[GraphProto],
-    initialization_bindings: Optional[AssignmentBindingType],
+    initialization: GraphProto | None,
+    initialization_bindings: AssignmentBindingType | None,
 ) -> TrainingInfoProto:
     training_info = TrainingInfoProto()
     training_info.algorithm.CopyFrom(algorithm)
@@ -1594,9 +1600,28 @@ def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
     Returns:
         TensorsProto's data_type
     """
-    return cast(
-        int,
-        mapping._NP_TYPE_TO_TENSOR_TYPE[np_dtype],
+    if np_dtype in mapping._NP_TYPE_TO_TENSOR_TYPE:
+        return cast(
+            int,
+            mapping._NP_TYPE_TO_TENSOR_TYPE[np_dtype],
+        )
+
+    if np.issubdtype(np_dtype, np.str_):
+        return TensorProto.STRING
+
+    if np_dtype in {
+        custom_np_types.bfloat16,
+        custom_np_types.float8e4m3fn,
+        custom_np_types.float8e4m3fnuz,
+        custom_np_types.float8e5m2,
+        custom_np_types.float8e5m2fnuz,
+        custom_np_types.int4,
+        custom_np_types.uint4,
+    }:
+        return custom_np_types.mapping_name_to_data_type[np_dtype.descr[0][0]]
+
+    raise ValueError(
+        f"Unable to convert type {np_dtype!r} into TensorProto element type."
     )