onnx 1.16.2__cp312-cp312-win_amd64.whl → 1.18.0__cp312-cp312-win_amd64.whl

onnx/hub.py

@@ -5,13 +5,16 @@
 
 This implements the python client for the ONNX model hub.
 """
+
+from __future__ import annotations
+
 import hashlib
 import json
 import os
 import sys
 from io import BytesIO
 from os.path import join
-from typing import IO, Any, Dict, List, Optional, Set, Tuple, cast
+from typing import IO, Any, cast
 from urllib.error import HTTPError
 from urllib.request import urlopen
 
@@ -38,7 +41,7 @@ class ModelInfo:
         opset: The opset version of the model.
     """
 
-    def __init__(self, raw_model_info: Dict[str, Any]) -> None:
+    def __init__(self, raw_model_info: dict[str, Any]) -> None:
         """Initializer.
 
         Args:
@@ -47,17 +50,17 @@ class ModelInfo:
         self.model = cast(str, raw_model_info["model"])
 
         self.model_path = cast(str, raw_model_info["model_path"])
-        self.metadata: Dict[str, Any] = cast(Dict[str, Any], raw_model_info["metadata"])
-        self.model_sha: Optional[str] = None
+        self.metadata: dict[str, Any] = cast(dict[str, Any], raw_model_info["metadata"])
+        self.model_sha: str | None = None
         if "model_sha" in self.metadata:
             self.model_sha = cast(str, self.metadata["model_sha"])
 
-        self.tags: Set[str] = set()
+        self.tags: set[str] = set()
         if "tags" in self.metadata:
-            self.tags = set(cast(List[str], self.metadata["tags"]))
+            self.tags = set(cast(list[str], self.metadata["tags"]))
 
         self.opset = cast(int, raw_model_info["opset_version"])
-        self.raw_model_info: Dict[str, Any] = raw_model_info
+        self.raw_model_info: dict[str, Any] = raw_model_info
 
     def __str__(self) -> str:
         return f"ModelInfo(model={self.model}, opset={self.opset}, path={self.model_path}, metadata={self.metadata})"
@@ -85,7 +88,7 @@ def get_dir() -> str:
     return _ONNX_HUB_DIR
 
 
-def _parse_repo_info(repo: str) -> Tuple[str, str, str]:
+def _parse_repo_info(repo: str) -> tuple[str, str, str]:
     """Gets the repo owner, name and ref from a repo specification string."""
     repo_owner = repo.split(":")[0].split("/")[0]
     repo_name = repo.split(":")[0].split("/")[1]
@@ -142,9 +145,9 @@ def _download_file(url: str, file_name: str) -> None:
 
 def list_models(
     repo: str = "onnx/models:main",
-    model: Optional[str] = None,
-    tags: Optional[List[str]] = None,
-) -> List[ModelInfo]:
+    model: str | None = None,
+    tags: list[str] | None = None,
+) -> list[ModelInfo]:
     """Gets the list of model info consistent with a given name and tags
 
     Args:
@@ -163,7 +166,7 @@ def list_models(
     manifest_url = base_url + "ONNX_HUB_MANIFEST.json"
     try:
         with urlopen(manifest_url) as response:
-            manifest: List[ModelInfo] = [
+            manifest: list[ModelInfo] = [
                 ModelInfo(info) for info in json.load(cast(IO[str], response))
             ]
     except HTTPError as e:
@@ -181,7 +184,7 @@ def list_models(
         return matching_models
 
     canonical_tags = {t.lower() for t in tags}
-    matching_info_list: List[ModelInfo] = []
+    matching_info_list: list[ModelInfo] = []
     for m in matching_models:
         model_tags = {t.lower() for t in m.tags}
         if len(canonical_tags.intersection(model_tags)) > 0:
@@ -190,7 +193,7 @@ def list_models(
 
 
 def get_model_info(
-    model: str, repo: str = "onnx/models:main", opset: Optional[int] = None
+    model: str, repo: str = "onnx/models:main", opset: int | None = None
 ) -> ModelInfo:
     """Gets the model info matching the given name and opset.
 
@@ -225,10 +228,10 @@ def get_model_info(
 def load(
     model: str,
     repo: str = "onnx/models:main",
-    opset: Optional[int] = None,
+    opset: int | None = None,
     force_reload: bool = False,
     silent: bool = False,
-) -> Optional[onnx.ModelProto]:
+) -> onnx.ModelProto | None:
     """Downloads a model by name from the onnx model hub.
 
     Args:
@@ -250,9 +253,9 @@ def load(
     selected_model = get_model_info(model, repo, opset)
     local_model_path_arr = selected_model.model_path.split("/")
     if selected_model.model_sha is not None:
-        local_model_path_arr[
-            -1
-        ] = f"{selected_model.model_sha}_{local_model_path_arr[-1]}"
+        local_model_path_arr[-1] = (
+            f"{selected_model.model_sha}_{local_model_path_arr[-1]}"
+        )
     local_model_path = join(_ONNX_HUB_DIR, os.sep.join(local_model_path_arr))
 
     if force_reload or not os.path.exists(local_model_path):
@@ -290,10 +293,10 @@ def load(
 def download_model_with_test_data(
     model: str,
     repo: str = "onnx/models:main",
-    opset: Optional[int] = None,
+    opset: int | None = None,
     force_reload: bool = False,
     silent: bool = False,
-) -> Optional[str]:
+) -> str | None:
     """Downloads a model along with test data by name from the onnx model hub and returns the directory to which the files have been extracted.
     Users are responsible for making sure the model comes from a trusted source, and the data is safe to be extracted.
 
@@ -322,9 +325,9 @@ def download_model_with_test_data(
     model_with_data_sha = selected_model.metadata["model_with_data_sha"]
 
     if model_with_data_sha is not None:
-        local_model_with_data_path_arr[
-            -1
-        ] = f"{model_with_data_sha}_{local_model_with_data_path_arr[-1]}"
+        local_model_with_data_path_arr[-1] = (
+            f"{model_with_data_sha}_{local_model_with_data_path_arr[-1]}"
+        )
     local_model_with_data_path = join(
         _ONNX_HUB_DIR, os.sep.join(local_model_with_data_path_arr)
     )
@@ -361,10 +364,8 @@ def download_model_with_test_data(
                 "download the model from the model hub."
             )
 
-    # FIXME: Avoid index manipulation with magic numbers,
-    # remove ".tar.gz"
     local_model_with_data_dir_path = local_model_with_data_path[
-        0 : len(local_model_with_data_path) - 7
+        0 : len(local_model_with_data_path) - len(".tar.gz")
     ]
     onnx.utils._extract_model_safe(
         local_model_with_data_path, local_model_with_data_dir_path
@@ -383,10 +384,10 @@ def load_composite_model(
     preprocessing_model: str,
     network_repo: str = "onnx/models:main",
     preprocessing_repo: str = "onnx/models:main",
-    opset: Optional[int] = None,
+    opset: int | None = None,
     force_reload: bool = False,
     silent: bool = False,
-) -> Optional[onnx.ModelProto]:
+) -> onnx.ModelProto | None:
     """Builds a composite model including data preprocessing by downloading a network and a preprocessing model
     and combine it into a single model
 
@@ -420,9 +421,9 @@ def load_composite_model(
     if network is None:
         raise RuntimeError(f"Could not load the network model: {network_model}")
 
-    all_domains: Set[str] = set()
-    domains_to_version_network: Dict[str, int] = {}
-    domains_to_version_preprocessing: Dict[str, int] = {}
+    all_domains: set[str] = set()
+    domains_to_version_network: dict[str, int] = {}
+    domains_to_version_preprocessing: dict[str, int] = {}
 
     for opset_import_entry in network.opset_import:
         domain = (

onnx/inliner/inliner.cc

@@ -15,9 +15,9 @@
 
 #include "onnx/common/assertions.h"
 #include "onnx/common/constants.h"
-#include "onnx/common/interned_strings.h"
 #include "onnx/common/proto_util.h"
 #include "onnx/common/visitor.h"
+#include "onnx/defs/parser.h"
 #include "onnx/shape_inference/attribute_binder.h"
 #include "onnx/shape_inference/implementation.h"
 #include "onnx/version_converter/convert.h"
@@ -66,7 +66,6 @@ struct OpsetMap : public OpsetMapBase {
   }
 
   bool Add(const google::protobuf::RepeatedPtrField<OperatorSetIdProto>& list) {
-    OpsetMapBase result;
     for (const auto& pair : list) {
       auto domain = NormalizeDomain(pair.domain());
       auto version = pair.version();
@@ -87,11 +86,11 @@ using RepeatedNodeProto = google::protobuf::RepeatedPtrField<NodeProto>;
 class NameGenerator : private Visitor {
  public:
   explicit NameGenerator(const GraphProto& graph) : index_(0) {
-    VisitGraph(graph);
+    NameGenerator::VisitGraph(graph);
   }
 
   explicit NameGenerator(const FunctionProto& function) : index_(0) {
-    VisitFunction(function);
+    NameGenerator::VisitFunction(function);
   }
 
   // Creates a new unique name, based on a suggested name, and adds it to the set
@@ -218,7 +217,7 @@ class InliningRenamer : private MutableVisitor {
     }
     for (; i < formals.size(); ++i) {
       std::string& formal = *formals.Mutable(i);
-      std::string rename_as = isOutput ? MakeUnique(formal) : std::string("");
+      std::string rename_as = isOutput ? MakeUnique(formal) : std::string();
       current_scope[formal] = rename_as;
       if (!rename_as.empty())
         formal = rename_as;
@@ -329,9 +328,9 @@ class ComputeInputs : private Visitor {
  public:
   std::vector<std::string> result;
 
-  ComputeInputs(const NodeProto& node) {
+  explicit ComputeInputs(const NodeProto& node) {
     result.reserve(node.input_size());
-    VisitNode(node);
+    ComputeInputs::VisitNode(node);
   }
 };
 
@@ -351,7 +350,7 @@ ConstNodeMap FindConstantNodes(const GraphProto& graph) {
   return result;
 }
 
-const TypeProto& GetType(const ModelProto& model, std::string var) {
+const TypeProto& GetType(const ModelProto& model, const std::string& var) {
   for (auto& vi : model.graph().value_info()) {
     if (vi.name() == var)
       return vi.type();
@@ -580,8 +579,7 @@ class VectorSet : public FunctionIdSet {
 std::unique_ptr<FunctionIdSet> ONNX_NAMESPACE::inliner::FunctionIdSet::Create(
     FunctionIdVector&& function_ids,
     bool invert) {
-  auto* p = new VectorSet(std::move(function_ids), invert);
-  return std::unique_ptr<FunctionIdSet>(p);
+  return std::make_unique<VectorSet>(std::move(function_ids), invert);
 }
 
 void InlineLocalFunctions(ModelProto& model, bool convert_version) {

onnx/mapping.py

@@ -1,9 +1,10 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import warnings
-from typing import Any, Dict, NamedTuple, Union, cast
+from typing import Any, NamedTuple, Union, cast
 
 import numpy as np
 
@@ -87,13 +88,16 @@ TENSOR_TYPE_MAP = {
     int(TensorProto.INT4): TensorDtypeMap(
         np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT4"
     ),
+    int(TensorProto.FLOAT4E2M1): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT4E2M1"
+    ),
 }
 
 
-class DeprecatedWarningDict(dict):  # type: ignore
+class DeprecatedWarningDict(dict):
     def __init__(
         self,
-        dictionary: Dict[int, Union[int, str, np.dtype]],
+        dictionary: dict[int, int | str | np.dtype],
         original_function: str,
         future_function: str = "",
     ) -> None:
@@ -109,7 +113,7 @@ class DeprecatedWarningDict(dict):  # type: ignore
             and self._future_function == other._future_function
         )
 
-    def __getitem__(self, key: Union[int, str, np.dtype]) -> Any:
+    def __getitem__(self, key: int | str | np.dtype) -> Any:
         if not self._future_function:
             warnings.warn(
                 str(
@@ -163,13 +167,14 @@ _NP_TYPE_TO_TENSOR_TYPE = {
         TensorProto.FLOAT8E5M2FNUZ,
         TensorProto.UINT4,
         TensorProto.INT4,
+        TensorProto.FLOAT4E2M1,
     )
 }
 
 # Currently native numpy does not support bfloat16 so TensorProto.BFLOAT16 is ignored for now
 # Numpy float32 array is only reversed to TensorProto.FLOAT
 NP_TYPE_TO_TENSOR_TYPE = DeprecatedWarningDict(
-    cast(Dict[int, Union[int, str, Any]], _NP_TYPE_TO_TENSOR_TYPE),
+    cast(dict[int, Union[int, str, Any]], _NP_TYPE_TO_TENSOR_TYPE),
     "NP_TYPE_TO_TENSOR_TYPE",
     "np_dtype_to_tensor_dtype",
 )
@@ -192,7 +197,7 @@ _STORAGE_TENSOR_TYPE_TO_FIELD = {
 }
 
 STORAGE_TENSOR_TYPE_TO_FIELD = DeprecatedWarningDict(
-    cast(Dict[int, Union[int, str, Any]], _STORAGE_TENSOR_TYPE_TO_FIELD),
+    cast(dict[int, Union[int, str, Any]], _STORAGE_TENSOR_TYPE_TO_FIELD),
     "STORAGE_TENSOR_TYPE_TO_FIELD",
 )
 

onnx/model_container.py

@@ -4,11 +4,12 @@
 """Implements function make_large_model to easily create and save models
 bigger than 2 Gb.
 """
+
 from __future__ import annotations
 
 import os
 import sys
-from typing import Any, Iterable
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -17,6 +18,9 @@ import onnx.external_data_helper as ext_data
 import onnx.helper
 import onnx.onnx_cpp2py_export.checker as c_checker
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 def _set_external_data(
     tensor: onnx.TensorProto,
@@ -79,7 +83,7 @@ class ModelContainer:
     No tensor is stored on disk until the user explicitly saves the model.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.model_proto_: onnx.ModelProto | None = None
         self.large_initializers: dict[str, np.ndarray] = {}
 
@@ -126,13 +130,13 @@ class ModelContainer:
                 )
         self.large_initializers = large_initializers
 
-    def check_large_initializers(self):
+    def check_large_initializers(self) -> None:
         for tensor in ext_data._get_all_tensors(self.model_proto):
             if not ext_data.uses_external_data(tensor):
                 continue
             prop: onnx.StringStringEntryProto | None = None
-            for ext in tensor.external_data:  # type: ignore[assignment]
-                if ext.key == "location":  # type: ignore[attr-defined]
+            for ext in tensor.external_data:
+                if ext.key == "location":
                     prop = ext
             if prop is None:
                 raise RuntimeError(
@@ -195,9 +199,9 @@ class ModelContainer:
             if not ext_data.uses_external_data(tensor):
                 continue
             prop: onnx.StringStringEntryProto | None = None
-            for ext in tensor.external_data:  # type: ignore[assignment]
-                if ext.key == "location":  # type: ignore[attr-defined]
-                    prop = ext  # type: ignore[assignment]
+            for ext in tensor.external_data:
+                if ext.key == "location":
+                    prop = ext
             if prop is None:
                 raise RuntimeError(
                     f"No location found for tensor name {tensor.name!r}."