onnx 1.15.0__cp311-cp311-win_amd64.whl → 1.16.1__cp311-cp311-win_amd64.whl

onnx/hub.py

@@ -27,8 +27,7 @@ else:
 
 
 class ModelInfo:
-    """
-    A class to represent a model's property and metadata in the ONNX Hub.
+    """A class to represent a model's property and metadata in the ONNX Hub.
     It extracts model name, path, sha, tags, etc. from the passed in raw_model_info dict.
 
     Attributes:
@@ -41,8 +40,9 @@ class ModelInfo:
     """
 
     def __init__(self, raw_model_info: Dict[str, Any]) -> None:
-        """
-        Parameters:
+        """Initializer.
+
+        Args:
             raw_model_info: A JSON dict containing the model info.
         """
         self.model = cast(str, raw_model_info["model"])
@@ -68,28 +68,26 @@ class ModelInfo:
 
 
 def set_dir(new_dir: str) -> None:
-    """
-    Sets the current ONNX hub cache location
+    """Sets the current ONNX hub cache location.
 
-    :param new_dir: location of new model hub cache
+    Args:
+        new_dir: Location of new model hub cache.
     """
     global _ONNX_HUB_DIR  # noqa: PLW0603
     _ONNX_HUB_DIR = new_dir
 
 
 def get_dir() -> str:
-    """
-    Gets the current ONNX hub cache location
+    """Gets the current ONNX hub cache location.
 
-    :return: The location of the ONNX hub model cache
+    Returns:
+        The location of the ONNX hub model cache.
     """
     return _ONNX_HUB_DIR
 
 
 def _parse_repo_info(repo: str) -> Tuple[str, str, str]:
-    """
-    Gets the repo owner, name and ref from a repo specification string.
-    """
+    """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]
     if ":" in repo:
@@ -100,8 +98,7 @@ def _parse_repo_info(repo: str) -> Tuple[str, str, str]:
 
 
 def _verify_repo_ref(repo: str) -> bool:
-    """
-    Verifies whether the given model repo can be trusted.
+    """Verifies whether the given model repo can be trusted.
     A model repo can be trusted if it matches onnx/models:main.
     """
     repo_owner, repo_name, repo_ref = _parse_repo_info(repo)
@@ -109,13 +106,16 @@ def _verify_repo_ref(repo: str) -> bool:
 
 
 def _get_base_url(repo: str, lfs: bool = False) -> str:
-    """
-    Gets the base github url from a repo specification string
+    """Gets the base github url from a repo specification string.
 
-    :param repo: The location of the model repo in format "user/repo[:branch]".
-        If no branch is found will default to "main"
-    :param lfs: whether the url is for downloading lfs models
-    :return: the base github url for downloading
+    Args:
+        repo: The location of the model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main".
+        lfs: Whether the url is for downloading lfs models.
+
+    Returns:
+        The base github url for downloading.
     """
     repo_owner, repo_name, repo_ref = _parse_repo_info(repo)
 
@@ -125,11 +125,11 @@ def _get_base_url(repo: str, lfs: bool = False) -> str:
 
 
 def _download_file(url: str, file_name: str) -> None:
-    """
-    Downloads the file with specifed file_name from the url
+    """Downloads the file with specified file_name from the url.
 
-    :param url: a url of download link
-    :param file_name: a specified file name for the downloaded file
+    Args:
+        url: A url of download link.
+        file_name: A specified file name for the downloaded file.
     """
     chunk_size = 16384  # 1024 * 16
     with urlopen(url) as response, open(file_name, "wb") as f:
@@ -146,14 +146,19 @@ def list_models(
     model: Optional[str] = None,
     tags: Optional[List[str]] = None,
 ) -> List[ModelInfo]:
-    """
-    Gets the list of model info consistent with a given name and tags
-
-    :param repo: The location of the model repo in format "user/repo[:branch]".
-        If no branch is found will default to "main"
-    :param model: The name of the model to search for. If `None`, will return all models with matching tags.
-    :param tags: A list of tags to filter models by. If `None`, will return all models with matching name.
-    :return: list of ModelInfo
+    """Gets the list of model info consistent with a given name and tags
+
+    Args:
+        repo: The location of the model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main"
+        model: The name of the model to search for. If `None`, will
+            return all models with matching tags.
+        tags: A list of tags to filter models by. If `None`, will return
+            all models with matching name.
+
+    Returns:
+        ``ModelInfo``s.
     """
     base_url = _get_base_url(repo)
     manifest_url = base_url + "ONNX_HUB_MANIFEST.json"
@@ -188,14 +193,19 @@ def list_models(
 def get_model_info(
     model: str, repo: str = "onnx/models:main", opset: Optional[int] = None
 ) -> ModelInfo:
-    """
-    Gets the model info matching the given name and opset.
-
-    :param model: The name of the onnx model in the manifest. This field is case-sensitive
-    :param repo: The location of the model repo in format "user/repo[:branch]".
-        If no branch is found will default to "main"
-    :param opset: The opset of the model to get. The default of `None` will return the model with largest opset.
-    :return: ModelInfo
+    """Gets the model info matching the given name and opset.
+
+    Args:
+        model: The name of the onnx model in the manifest. This field is
+            case-sensitive
+        repo: The location of the model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main"
+        opset: The opset of the model to get. The default of `None` will
+            return the model with largest opset.
+
+    Returns:
+        ``ModelInfo``.
     """
     matching_models = list_models(repo, model)
     if not matching_models:
@@ -220,16 +230,23 @@ def load(
     force_reload: bool = False,
     silent: bool = False,
 ) -> Optional[onnx.ModelProto]:
-    """
-    Downloads a model by name from the onnx model hub
-
-    :param model: The name of the onnx model in the manifest. This field is case-sensitive
-    :param repo: The location of the model repo in format "user/repo[:branch]".
-        If no branch is found will default to "main"
-    :param opset: The opset of the model to download. The default of `None` automatically chooses the largest opset
-    :param force_reload: Whether to force the model to re-download even if its already found in the cache
-    :param silent: Whether to suppress the warning message if the repo is not trusted.
-    :return: ModelProto or None
+    """Downloads a model by name from the onnx model hub.
+
+    Args:
+        model: The name of the onnx model in the manifest. This field is
+            case-sensitive
+        repo: The location of the model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main"
+        opset: The opset of the model to download. The default of `None`
+            automatically chooses the largest opset
+        force_reload: Whether to force the model to re-download even if
+            its already found in the cache
+        silent: Whether to suppress the warning message if the repo is
+            not trusted.
+
+    Returns:
+        ModelProto or None
     """
     selected_model = get_model_info(model, repo, opset)
     local_model_path_arr = selected_model.model_path.split("/")
@@ -278,16 +295,23 @@ def download_model_with_test_data(
     force_reload: bool = False,
     silent: bool = False,
 ) -> Optional[str]:
-    """
-    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.
-
-    :param model: The name of the onnx model in the manifest. This field is case-sensitive
-    :param repo: The location of the model repo in format "user/repo[:branch]".
-        If no branch is found will default to "main"
-    :param opset: The opset of the model to download. The default of `None` automatically chooses the largest opset
-    :param force_reload: Whether to force the model to re-download even if its already found in the cache
-    :param silent: Whether to suppress the warning message if the repo is not trusted.
-    :return: str or 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.
+
+    Args:
+        model: The name of the onnx model in the manifest. This field is
+            case-sensitive
+        repo: The location of the model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main"
+        opset: The opset of the model to download. The default of `None`
+            automatically chooses the largest opset
+        force_reload: Whether to force the model to re-download even if
+            its already found in the cache
+        silent: Whether to suppress the warning message if the repo is
+            not trusted.
+
+    Returns:
+        str or None
     """
     selected_model = get_model_info(model, repo, opset)
 
@@ -361,17 +385,27 @@ def load_composite_model(
     force_reload: bool = False,
     silent: bool = False,
 ) -> Optional[onnx.ModelProto]:
-    """
-    Builds a composite model including data preprocessing by downloading a network and a preprocessing model
+    """Builds a composite model including data preprocessing by downloading a network and a preprocessing model
     and combine it into a single model
 
-    :param model: The name of the onnx model in the manifest. This field is case-sensitive
-    :param repo: The location of the model repo in format "user/repo[:branch]".
-        If no branch is found will default to "main"
-    :param opset: The opset of the model to download. The default of `None` automatically chooses the largest opset
-    :param force_reload: Whether to force the model to re-download even if its already found in the cache
-    :param silent: Whether to suppress the warning message if the repo is not trusted.
-    :return: ModelProto or None
+    Args:
+        network_model: The name of the onnx model in the manifest.
+        preprocessing_model: The name of the preprocessing model.
+        network_repo: The location of the model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main"
+        preprocessing_repo: The location of the proprocessing model repo in format
+            "user/repo[:branch]". If no branch is found will default to
+            "main"
+        opset: The opset of the model to download. The default of `None`
+            automatically chooses the largest opset
+        force_reload: Whether to force the model to re-download even if
+            its already found in the cache
+        silent: Whether to suppress the warning message if the repo is
+            not trusted.
+
+    Returns:
+        ModelProto or None
     """
     preprocessing = load(
         preprocessing_model, preprocessing_repo, opset, force_reload, silent

onnx/inliner/inliner.cc

@@ -16,6 +16,7 @@
 #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/shape_inference/attribute_binder.h"
 #include "onnx/shape_inference/implementation.h"
@@ -28,27 +29,6 @@ namespace { // internal/private API
 
 using namespace internal;
 
-// Function lookup function. Returns true iff lookup is successful, in which case
-// the found FunctionProto is copied into *return_value. The opset version is not
-// a parameter since it is determined by the domain for a given model.
-using FunctionResolver =
-    std::function<bool(const std::string& domain, const std::string& op, FunctionProto* return_value)>;
-
-// We use a string of the form "domain::name" as the key for a function id.
-using FunctionIdKey = std::string;
-
-FunctionIdKey GetFunctionId(const std::string& domain, const std::string& op) {
-  return NormalizeDomain(domain) + "::" + op;
-}
-
-FunctionIdKey GetFunctionId(const FunctionProto& function) {
-  return GetFunctionId(function.domain(), function.name());
-}
-
-FunctionIdKey GetCalleeId(const NodeProto& node) {
-  return GetFunctionId(node.domain(), node.op_type());
-}
-
 using OpsetMapBase = std::unordered_map<std::string, int64_t>;
 
 // A representation of the opset versions required by a model or a function.
@@ -287,6 +267,8 @@ class InliningRenamer : private MutableVisitor {
     renamer.Bind<true>(*callee.mutable_output(), callnode.output());
 
     renamer.VisitFunction(&callee);
+    for (auto& v : *callee.mutable_value_info())
+      renamer.LookupOrRename(*v.mutable_name(), false);
   }
 };
 
@@ -441,17 +423,43 @@ void ConvertVersion(ModelProto& model, const NodeProto& call_node, FunctionProto
 }
 
 constexpr int64_t kNoConversion = -1;
-using FunctionMap = std::unordered_map<FunctionIdKey, std::pair<const FunctionProto*, int64_t>>;
+using FunctionMap = std::unordered_map<FunctionImplId, std::pair<const FunctionProto*, int64_t>>;
 
 using NodeList = google::protobuf::RepeatedPtrField<NodeProto>;
 
-// Shared utility used for inlining into either a GraphProto or a FunctionProto.
-void InlineFunctions(NodeList& nodes, const FunctionMap& map, NameGenerator& name_generator, ModelProto* model) {
+/** Utility function used for inlining into a GraphProto.
+ * @param graph Mutable graph
+ * @param map Map from function-id to function for functions to be inlined
+ * @param name_generator Name generator for generating unique names for inlined variables
+ * @param model If non-null, the model being inlined into. Used for version conversion.
+ * @param inline_count Mutable counter for number of inlined calls. Used for name generation.
+ */
+void InlineFunctions(
+    GraphProto& graph,
+    const FunctionMap& map,
+    NameGenerator& name_generator,
+    ModelProto* model,
+    int& inline_count);
+
+/** Shared utility function used for inlining into either a GraphProto or a FunctionProto.
+ * @param nodes Mutable list of nodes (of function or graph)
+ * @param value_infos Mutable list of value_infos (of function or graph)
+ * @param map Map from function-id to function for functions to be inlined
+ * @param name_generator Name generator for generating unique names for inlined variables
+ * @param model If non-null, the model being inlined into. Used for version conversion.
+ * @param inline_count Mutable counter for number of inlined calls. Used for name generation.
+ */
+void InlineFunctions(
+    NodeList& nodes,
+    ValueInfoList& value_infos,
+    const FunctionMap& map,
+    NameGenerator& name_generator,
+    ModelProto* model,
+    int& inline_count) {
   NodeList original_nodes;
   // Move all nodes into original_nodes
   original_nodes.Swap(&nodes);
 
-  int inline_count = 0;
   std::function<void(NodeProto & node)> append_node = [&](NodeProto& node) {
     FunctionProto callee;
     auto iter = map.find(GetCalleeId(node));
@@ -472,6 +480,17 @@ void InlineFunctions(NodeList& nodes, const FunctionMap& map, NameGenerator& nam
             callee.name().c_str());
         ConvertVersion(*model, node, callee, target_version);
       }
+      std::unordered_set<std::string> actual_parameters;
+      for (const auto& x : node.input())
+        actual_parameters.insert(x);
+      for (const auto& x : node.output())
+        actual_parameters.insert(x);
+      // Append valueinfos of called function
+      for (auto& callee_vi : callee.value_info()) {
+        if (actual_parameters.count(callee_vi.name()) == 0) {
+          *value_infos.Add() = callee_vi;
+        }
+      }
       // Append nodes of called function
       for (auto& callee_node : *callee.mutable_node())
         append_node(callee_node);
@@ -479,6 +498,16 @@ void InlineFunctions(NodeList& nodes, const FunctionMap& map, NameGenerator& nam
       // Append node without inlining.
       // TODO: use std::move instead of copying. Use of move doesn't seem to work with
       // protobuf in some platforms/settings. [nodes->Add(std::move(node));]
+
+      for (auto& attr : *node.mutable_attribute()) {
+        if (attr.has_g()) {
+          InlineFunctions(*attr.mutable_g(), map, name_generator, model, inline_count);
+        }
+        for (auto& g : *attr.mutable_graphs()) {
+          InlineFunctions(g, map, name_generator, model, inline_count);
+        }
+      }
+
       *nodes.Add() = node;
     }
   };
@@ -487,16 +516,45 @@ void InlineFunctions(NodeList& nodes, const FunctionMap& map, NameGenerator& nam
   }
 }
 
+/** Utility function used for inlining into a GraphProto.
+ * @param graph Mutable graph
+ * @param map Map from function-id to function for functions to be inlined
+ * @param name_generator Name generator for generating unique names for inlined variables
+ * @param model If non-null, the model being inlined into. Used for version conversion.
+ * @param inline_count Mutable counter for number of inlined calls. Used for name generation.
+ */
+void InlineFunctions(
+    GraphProto& graph,
+    const FunctionMap& map,
+    NameGenerator& name_generator,
+    ModelProto* model,
+    int& inline_count) {
+  auto* nodes = graph.mutable_node();
+  auto* value_infos = graph.mutable_value_info();
+  InlineFunctions(*nodes, *value_infos, map, name_generator, model, inline_count);
+}
+
+/** Utility function used for inlining into a ModelProto.
+ * @param model Mutable model
+ * @param map Map from function-id to function for functions to be inlined
+ */
 void InlineFunctions(ModelProto& model, FunctionMap& map) {
+  int inline_count = 0;
   auto* graph = model.mutable_graph();
   NameGenerator name_generator(*graph);
   auto* nodes = graph->mutable_node();
-  InlineFunctions(*nodes, map, name_generator, &model);
+  auto* value_infos = graph->mutable_value_info();
+  InlineFunctions(*nodes, *value_infos, map, name_generator, &model, inline_count);
 }
 
+/** Utility function used for inlining into a FunctionProto.
+ * @param function Mutable function
+ * @param map Map from function-id to function for functions to be inlined
+ */
 void InlineFunctions(FunctionProto& function, FunctionMap& map) {
+  int inline_count = 0;
   NameGenerator name_generator(function);
-  InlineFunctions(*function.mutable_node(), map, name_generator, nullptr);
+  InlineFunctions(*function.mutable_node(), *function.mutable_value_info(), map, name_generator, nullptr, inline_count);
 }
 
 class VectorSet : public FunctionIdSet {
@@ -546,7 +604,7 @@ void InlineLocalFunctions(ModelProto& model, bool convert_version) {
       mismatches.erase(iter);
     }
     if (mismatches.empty()) {
-      map[GetFunctionId(function)] = std::pair<const FunctionProto*, int64_t>(&function, target_onnx_version);
+      map[GetFunctionImplId(function)] = std::pair<const FunctionProto*, int64_t>(&function, target_onnx_version);
     }
   }
 
@@ -556,7 +614,7 @@ void InlineLocalFunctions(ModelProto& model, bool convert_version) {
   // opset version. They need to be handled some other way, eg., using a version-adapter.
   auto* local_functions = model.mutable_functions();
   for (auto it = local_functions->begin(); it != local_functions->end();) {
-    if (map.count(GetFunctionId(*it)) > 0)
+    if (map.count(GetFunctionImplId(*it)) > 0)
       it = local_functions->erase(it);
     else
       ++it;
@@ -576,7 +634,7 @@ void InlineSelectedFunctions(ModelProto& model, const FunctionIdSet& to_inline)
     if (!model_imports.Add(function))
       ONNX_THROW("Model has functions with incompatible opset versions.");
     if (to_inline.Contains(function.domain(), function.name())) {
-      map[GetFunctionId(function)] = std::pair<const FunctionProto*, int64_t>(&function, kNoConversion);
+      map[GetFunctionImplId(function)] = std::pair<const FunctionProto*, int64_t>(&function, kNoConversion);
     } else {
       non_inlined_functions.push_back(&function);
     }
@@ -591,7 +649,7 @@ void InlineSelectedFunctions(ModelProto& model, const FunctionIdSet& to_inline)
   // Remove all inlined model-local functions.
   auto* local_functions = model.mutable_functions();
   for (auto it = local_functions->begin(); it != local_functions->end();) {
-    if (map.count(GetFunctionId(*it)) > 0)
+    if (map.count(GetFunctionImplId(*it)) > 0)
       it = local_functions->erase(it);
     else
       ++it;

onnx/inliner/inliner.h

@@ -14,6 +14,11 @@
 namespace ONNX_NAMESPACE {
 namespace inliner {
 
+// IR version 10 introduces overloaded function names. The following APIs to specify
+// functions to be inlined currently allow only specifying (domain, name). Thus,
+// either all overloads of a function are inlined or none.
+// The older-style ids are used below for backward compatibility.
+
 // A FunctionId is a pair of strings (domain, function name).
 using FunctionId = std::pair<std::string, std::string>;
 

onnx/inliner.py

@@ -17,6 +17,7 @@ def inline_local_functions(
         model: an ONNX ModelProto
         convert_version: if true, try to apply automatic version-conversion to functions requiring a
             different (ONNX) opset version from the model.
+
     Returns:
         ModelProto with all calls to model-local functions inlined (recursively)
     """
@@ -37,6 +38,7 @@ def inline_selected_functions(
             element is a tuple of (function domain, function name).
         exclude: if true, inlines all functions except those specified in function_ids.
            if false, inlines all functions specified in function_ids.
+
     Returns:
         ModelProto with all calls to model-local functions inlined (recursively)
     """

onnx/mapping.py

@@ -80,6 +80,13 @@ TENSOR_TYPE_MAP = {
     int(TensorProto.FLOAT8E5M2FNUZ): TensorDtypeMap(
         np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2FNUZ"
     ),
+    # Native numpy does not support uint4/int4 so now use uint8/int8 for these types.
+    int(TensorProto.UINT4): TensorDtypeMap(
+        np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.UINT4"
+    ),
+    int(TensorProto.INT4): TensorDtypeMap(
+        np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT4"
+    ),
 }
 
 
@@ -154,6 +161,8 @@ _NP_TYPE_TO_TENSOR_TYPE = {
         TensorProto.FLOAT8E4M3FNUZ,
         TensorProto.FLOAT8E5M2,
         TensorProto.FLOAT8E5M2FNUZ,
+        TensorProto.UINT4,
+        TensorProto.INT4,
     )
 }