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

onnx/reference/reference_evaluator.py

@@ -1,9 +1,10 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 from io import BytesIO
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 
@@ -182,7 +183,7 @@ class ReferenceEvaluator:
 
         An operator may have a reference implementation such as `CastLike`
         and still be defined as a function. By default, the reference implementation
-        is used. This behaviour can be changed by adding a class to the list
+        is used. This behavior can be changed by adding a class to the list
         of overwritten operators. It must inherit from :class:`OpRunExpand`.
 
         ::
@@ -199,15 +200,15 @@ class ReferenceEvaluator:
             implementation a schema may define.
     """
 
-    def __init__(  # type: ignore
+    def __init__(
         self,
         proto: Any,
-        opsets: Optional[Dict[str, int]] = None,
-        functions: Optional[List[Union["ReferenceEvaluator", FunctionProto]]] = None,  # type: ignore
+        opsets: dict[str, int] | None = None,
+        functions: list[ReferenceEvaluator | FunctionProto] | None = None,  # type: ignore
         verbose: int = 0,
-        new_ops: Optional[List[OpRun]] = None,
+        new_ops: list[type[OpRun]] | None = None,
         optimized: bool = True,
-    ):
+    ) -> None:
         if optimized:
             if new_ops is None:
                 new_ops = optimized_operators.copy()
@@ -220,7 +221,7 @@ class ReferenceEvaluator:
         self.input_types_ = None
 
         if isinstance(proto, ModelContainer):
-            self.container_ = proto
+            self.container_: ModelContainer | None = proto
             proto = self.container_.model_proto
         else:
             self.container_ = None
@@ -231,8 +232,8 @@ class ReferenceEvaluator:
         elif isinstance(proto, bytes):
             proto = load(BytesIO(proto))
         self.proto_ = proto
-        self.functions_: Dict[Tuple[str, str], ReferenceEvaluator] = {}
-        self.attributes_: List[str] = []
+        self.functions_: dict[tuple[str, str], ReferenceEvaluator] = {}
+        self.attributes_: list[str] = []
         if isinstance(proto, ModelProto):
             self.onnx_graph_ = proto.graph
             self.opsets_ = {d.domain: d.version for d in proto.opset_import}
@@ -244,7 +245,7 @@ class ReferenceEvaluator:
         elif isinstance(proto, GraphProto):
             self.onnx_graph_ = proto
             if not isinstance(opsets, dict):
-                raise ValueError("opsets must be a dictionary if proto is GraphProto.")
+                raise TypeError("opsets must be a dictionary if proto is GraphProto.")
             self.opsets_ = opsets
         elif isinstance(proto, FunctionProto):
             self.onnx_graph_ = None  # type: ignore
@@ -293,7 +294,7 @@ class ReferenceEvaluator:
                 else:
                     raise TypeError(f"Unexpected type {type(f)!r} for a function.")
         self.verbose = verbose
-        self.new_ops_: Dict[Tuple[str, str], OpRun] = {}
+        self.new_ops_: dict[tuple[str, str], type[OpRun]] = {}
         if new_ops is not None:
             for cl in new_ops:
                 if not hasattr(cl, "op_domain"):
@@ -309,7 +310,7 @@ class ReferenceEvaluator:
                 self.new_ops_[key] = cl
         self._init()
 
-    def retrieve_external_data(self, initializer: TensorProto) -> np.array:
+    def retrieve_external_data(self, initializer: TensorProto) -> np.ndarray:
         """Returns a tensor saved as external."""
         info = ExternalDataInfo(initializer)
         location = info.location
@@ -342,7 +343,7 @@ class ReferenceEvaluator:
             return ", ".join(map(self._log_arg, a))
         return a
 
-    def _log(self, level: int, pattern: str, *args: List[Any]) -> None:
+    def _log(self, level: int, pattern: str, *args: list[Any]) -> None:
         if level < self.verbose:
             new_args = [self._log_arg(a) for a in args]
             print(pattern % tuple(new_args))
@@ -453,7 +454,7 @@ class ReferenceEvaluator:
             self.rt_nodes_.append(inst)
 
     def _load_impl(  # noqa: PLR0911
-        self, node: NodeProto, input_types: Optional[TypeProto] = None
+        self, node: NodeProto, input_types: TypeProto | None = None
     ) -> Any:
         """Loads the implementation for a specified runtime."""
         if node.domain not in self.opsets:
@@ -493,7 +494,7 @@ class ReferenceEvaluator:
                     node.op_type,
                     version,
                     node=node,
-                    input_types=input_types,  # type: ignore[arg-type]
+                    input_types=input_types,
                     expand=expand,
                     evaluator_cls=self.__class__,
                 )
@@ -507,23 +508,17 @@ class ReferenceEvaluator:
         if node.domain == "ai.onnx.preview.training":
             from onnx.reference.ops.aionnx_preview_training import load_op as load_op_pt
 
-            return load_op_pt(
-                node.domain, node.op_type, version, evaluator_cls=self.__class__
-            )
+            return load_op_pt(node.domain, node.op_type, version)
 
         if node.domain == "experimental":
             from onnx.reference.ops.experimental import load_op as load_op_exp
 
-            return load_op_exp(
-                node.domain, node.op_type, version, evaluator_cls=self.__class__
-            )
+            return load_op_exp(node.domain, node.op_type, version)
 
         if node.domain == "ai.onnx.ml":
             from onnx.reference.ops.aionnxml import load_op as load_op_ml
 
-            return load_op_ml(
-                node.domain, node.op_type, version, evaluator_cls=self.__class__
-            )
+            return load_op_ml(node.domain, node.op_type, version)
 
         # It has to be a function.
         if key in self.functions_:
@@ -542,7 +537,13 @@ class ReferenceEvaluator:
             f"is unknown, known functions: {sorted(self.functions_)}."
         )
 
-    def run(self, output_names, feed_inputs: Dict[str, Any], attributes: Optional[Dict[str, Any]] = None):  # type: ignore
+    def run(
+        self,
+        output_names,
+        feed_inputs: dict[str, Any],
+        attributes: dict[str, Any] | None = None,
+        intermediate: bool = False,
+    ) -> dict[str, Any] | list[Any]:  # type: ignore
         """Executes the onnx model.
 
         Args:
@@ -550,9 +551,13 @@ class ReferenceEvaluator:
             feed_inputs: dictionary `{ input name: input value }`
             attributes: attributes value if the instance runs a
                 FunctionProto
+            intermediate: if True, the function returns all the results,
+                final ones and intermediates one in a same dictionary,
+                if False, only the final results are returned in a list
 
         Returns:
-            list of requested outputs
+            list of requested outputs if intermediate is False,
+            named results in a dictionary otherwise
         """
         if output_names is None:
             output_names = self.output_names
@@ -591,6 +596,9 @@ class ReferenceEvaluator:
                 results[name] = value
 
         # return the results
+        if intermediate:
+            return results
+
         for name in output_names:
             if name not in results:
                 raise RuntimeError(

onnx/serialization.py

@@ -11,7 +11,7 @@ __all__ = [
 ]
 
 import typing
-from typing import Any, Collection, Optional, Protocol, TypeVar
+from typing import Any, Optional, Protocol, TypeVar
 
 import google.protobuf.json_format
 import google.protobuf.message
@@ -19,6 +19,9 @@ import google.protobuf.text_format
 
 import onnx
 
+if typing.TYPE_CHECKING:
+    from collections.abc import Collection
+
 _Proto = TypeVar("_Proto", bound=google.protobuf.message.Message)
 # Encoding used for serializing and deserializing text files
 _ENCODING = "utf-8"
@@ -171,7 +174,7 @@ class _TextualSerializer(ProtoSerializer):
     """Serialize and deserialize the ONNX textual representation."""
 
     supported_format = "onnxtxt"
-    file_extensions = frozenset({".onnxtxt"})
+    file_extensions = frozenset({".onnxtxt", ".onnxtext"})
 
     def serialize_proto(self, proto: _Proto) -> bytes:
         text = onnx.printer.to_text(proto)  # type: ignore[arg-type]

onnx/shape_inference/attribute_binder.h

@@ -16,7 +16,7 @@ using AttributeMap = std::unordered_map<std::string, const AttributeProto*>;
 
 class AttributeBinder : public MutableVisitor {
  public:
-  AttributeBinder(const AttributeMap& attr_map) : attr_map_(attr_map) {}
+  explicit AttributeBinder(const AttributeMap& attr_map) : attr_map_(attr_map) {}
 
   // Binding a formal attribute-parameter to a value may, as a special case, also
   // remove the attribute from the list of attributes of a node (when the attribute

onnx/shape_inference/implementation.cc

@@ -16,7 +16,6 @@
 #include "onnx/common/common.h"
 #include "onnx/common/file_utils.h"
 #include "onnx/defs/data_type_utils.h"
-#include "onnx/proto_utils.h"
 #include "onnx/shape_inference/attribute_binder.h"
 #include "onnx/string_utils.h"
 
@@ -73,7 +72,7 @@ inline bool IsOnnxDomainOp(const NodeProto& node, const std::string& op_type) {
 } // namespace
 
 template <class T>
-void CheckTensorShapesAndTypes(const T& inferred_type, const T& existing_type) {
+static void CheckTensorShapesAndTypes(const T& inferred_type, const T& existing_type) {
   if (inferred_type.elem_type() != TensorProto::UNDEFINED && existing_type.elem_type() != TensorProto::UNDEFINED &&
       existing_type.elem_type() != inferred_type.elem_type()) {
     std::stringstream ss;
@@ -250,20 +249,20 @@ void MaterializeSymbolicShape(TypeProto* inferred_type, SymbolTable& symbol_tabl
   }
 }
 
-std::string GetFunctionIdentifier(const FunctionProto& function) {
+static std::string GetFunctionIdentifier(const FunctionProto& function) {
   // Note: Models with IR version < 10 do not have the overload attribute.
   // However, that will be mapped to an empty identifier.
-  std::string overload = function.overload();
+  const std::string& overload = function.overload();
   if (overload.empty()) {
     return function.domain() + ":" + function.name();
   }
   return function.domain() + ":" + function.name() + ":" + overload;
 }
 
-std::string GetFunctionIdentifier(const NodeProto& node) {
+static std::string GetFunctionIdentifier(const NodeProto& node) {
   // Note: Models with IR version < 10 do not have the overload attribute.
   // However, that will be mapped to an empty identifier.
-  std::string overload = node.overload();
+  const std::string& overload = node.overload();
   if (overload.empty()) {
     return node.domain() + ":" + node.op_type();
   }
@@ -304,7 +303,7 @@ class InferredTypes {
 };
 
 // Initialize a DataValueMap for a called function from the DataValueMap of the caller
-void BindValuesOnCall(
+static void BindValuesOnCall(
     const DataValueMap& caller_map,
     const NodeProto& caller,
     DataValueMap& callee_map,
@@ -323,7 +322,7 @@ void BindValuesOnCall(
 }
 
 // Update a DataValueMap for a calling function from the DataValueMap of the callee
-void BindValuesOnReturn(
+static void BindValuesOnReturn(
     const DataValueMap& callee_map,
     const FunctionProto& callee,
     DataValueMap& caller_map,
@@ -482,7 +481,7 @@ class ShapeInferenceImplBase {
         if (options.check_type) {
           schema->CheckInputOutputType(ctx);
         }
-      } else if (model_local_functions_map.size() > 0) {
+      } else if (!model_local_functions_map.empty()) {
         auto iter = model_local_functions_map.find(GetFunctionIdentifier(n));
         if (iter != model_local_functions_map.end()) {
           ProcessCall(n, *(iter->second), ctx);
@@ -787,7 +786,7 @@ static void InferShapesImpl(
 
 // Either ModelProto or FunctionProto
 template <class T>
-std::unordered_map<std::string, int> GetOpsetImportsFromProto(const T& proto) {
+static std::unordered_map<std::string, int> GetOpsetImportsFromProto(const T& proto) {
   std::unordered_map<std::string, int> opset_imports;
   for (const auto& opset_import : proto.opset_import()) {
     opset_imports[opset_import.domain()] = static_cast<int>(opset_import.version());
@@ -906,13 +905,13 @@ struct FunctionInferenceContext : public InferenceContext {
       const std::vector<TypeProto>& input_types,
       const std::vector<AttributeProto>& attributes,
       const ShapeInferenceOptions& options)
-      : input_types_(input_types), options_(options) {
+      : input_types_(input_types), options_(options), func_proto_(&func_proto) {
     for (const auto& attr : attributes) {
       attributesByName_[attr.name()] = &attr;
     }
     auto num_outputs = func_proto.output_size();
     for (int i = 0; i < num_outputs; i++) {
-      output_types_.push_back(TypeProto());
+      output_types_.emplace_back();
     }
   }
 
@@ -971,11 +970,25 @@ struct FunctionInferenceContext : public InferenceContext {
     return std::move(output_types_);
   }
 
+  std::string getDisplayName() const override {
+    if (func_proto_ == nullptr)
+      return "";
+    if (func_proto_->domain().empty()) {
+      if (func_proto_->name().empty())
+        return "";
+      return MakeString("function ", func_proto_->name());
+    }
+    if (func_proto_->name().empty())
+      return MakeString("function [", func_proto_->domain(), "]");
+    return MakeString("function ", func_proto_->name(), "[", func_proto_->domain(), "]");
+  }
+
  private:
   const std::vector<TypeProto>& input_types_;
   std::vector<TypeProto> output_types_;
   std::unordered_map<std::string, const AttributeProto*> attributesByName_;
   ShapeInferenceOptions options_;
+  const FunctionProto* func_proto_;
 };
 
 std::vector<TypeProto> InferFunctionOutputTypes(

onnx/shape_inference/implementation.h

@@ -4,7 +4,6 @@
 
 #pragma once
 
-#include <map>
 #include <memory>
 #include <string>
 #include <unordered_map>
@@ -12,9 +11,8 @@
 #include <utility>
 #include <vector>
 
-#include "onnx/defs/function.h"
 #include "onnx/defs/schema.h"
-#include "onnx/proto_utils.h"
+#include "onnx/defs/tensor_proto_util.h"
 #include "onnx/string_utils.h"
 
 namespace ONNX_NAMESPACE {
@@ -31,7 +29,7 @@ using DataValueMap = std::unordered_map<std::string, TensorShapeProto>;
 
 class SymbolTableImpl : public SymbolTable {
  public:
-  SymbolTableImpl() : index_(0) {}
+  SymbolTableImpl() = default;
 
   void addFromGraph(const GraphProto& g) override {
     AddExistingSymbolicDims(g.input());
@@ -50,7 +48,7 @@ class SymbolTableImpl : public SymbolTable {
   }
 
  private:
-  unsigned int index_;
+  unsigned int index_{0};
   std::unordered_set<std::string> existing_symbols;
 
   // TypeProto_Tensor or TypeProto_SparseTensor
@@ -98,14 +96,14 @@ class SymbolTableImpl : public SymbolTable {
 struct GraphInferenceContext {
   GraphInferenceContext(
       const std::unordered_map<std::string, TypeProto*>& outer_scope_value_types_by_name_in,
-      const std::unordered_map<std::string, int> opset_imports_in,
+      std::unordered_map<std::string, int> opset_imports_in,
       SymbolTable* symbol_table_in = nullptr,
       const ModelLocalFunctionsMap& model_local_functions_in = {},
       const ISchemaRegistry* schema_registry_in = OpSchemaRegistry::Instance(),
       DataValueMap* generated_shape_data_by_name_in = nullptr,
       const int ir_version_in = IR_VERSION)
       : outer_scope_value_types_by_name{&outer_scope_value_types_by_name_in},
-        opset_imports{opset_imports_in},
+        opset_imports{std::move(opset_imports_in)},
         symbol_table{symbol_table_in},
         model_local_functions{model_local_functions_in},
         schema_registry{schema_registry_in},
@@ -146,7 +144,7 @@ struct InferenceContextImpl : public InferenceContext {
       const ShapeInferenceOptions& options,
       DataValueMap* generatedShapeData = nullptr,
       GraphInferenceContext* graphInferenceContext = nullptr)
-      : graphInferenceContext_{graphInferenceContext}, options_(options) {
+      : graphInferenceContext_{graphInferenceContext}, options_(options), node_(&n) {
     for (auto& attr : *n.mutable_attribute()) {
       attributesByName_[attr.name()] = &attr;
       if (attr.has_g()) {
@@ -265,9 +263,8 @@ struct InferenceContextImpl : public InferenceContext {
         fail_type_inference("Attribute ", attr_name, " does not contain a graph.");
       }
 
-      std::unique_ptr<GraphInferencer> new_inferencer{
-          new GraphInferencerImpl(*attrNameToGraphProto->second, *graphInferenceContext_, options_)};
-
+      auto new_inferencer =
+          std::make_unique<GraphInferencerImpl>(*attrNameToGraphProto->second, *graphInferenceContext_, options_);
       inferencer = new_inferencer.get();
       graphAttributeInferencers_.emplace(attr_name, std::move(new_inferencer));
     } else {
@@ -277,6 +274,19 @@ struct InferenceContextImpl : public InferenceContext {
     return inferencer;
   }
 
+  std::string getDisplayName() const override {
+    if (node_ == nullptr)
+      return "";
+    if (node_->domain().empty()) {
+      if (node_->name().empty())
+        return MakeString("node ", node_->op_type());
+      return MakeString("node ", node_->op_type(), " (", node_->name(), ")");
+    }
+    if (node_->name().empty())
+      return MakeString("node ", node_->op_type(), "[", node_->domain(), "]");
+    return MakeString("node ", node_->op_type(), "[", node_->domain(), "]", " (", node_->name(), ")");
+  }
+
   std::vector<const TensorProto*> allInputData_;
   std::vector<const SparseTensorProto*> allInputSparseData_;
   std::vector<const TensorShapeProto*> allShapeInputData_;
@@ -289,6 +299,7 @@ struct InferenceContextImpl : public InferenceContext {
   // mutable as internal cache of GraphInferencer instances
   mutable std::unordered_map<std::string, std::unique_ptr<GraphInferencer>> graphAttributeInferencers_;
   ShapeInferenceOptions options_;
+  NodeProto* node_;
 };
 
 struct DataPropagationContextImpl : public DataPropagationContext {
@@ -402,6 +413,27 @@ struct DataPropagationContextImpl : public DataPropagationContext {
         return &(result.first->second);
       }
     }
+
+    // If X has a known rank N, then X's value can be represented as an array
+    // of N unknown values (represented as a TensorShapeProto).
+    const TypeProto* type = getInputType(index);
+    if ((type != nullptr) && (type->has_tensor_type())) {
+      auto& tensor_type = type->tensor_type();
+      if (tensor_type.has_shape()) {
+        const TensorShapeProto& shape = tensor_type.shape();
+        if ((shape.dim_size() == 1) && (shape.dim(0).has_dim_value())) {
+          TensorShapeProto tsp;
+          int64_t dim_value = shape.dim(0).dim_value();
+          for (int64_t i = 0; i < dim_value; ++i) {
+            tsp.add_dim();
+          }
+          auto result = generatedShapeData_.insert({input_name, std::move(tsp)});
+          if (result.second) {
+            return &(result.first->second);
+          }
+        }
+      }
+    }
     return nullptr;
   }
 
@@ -445,24 +477,24 @@ void mergeShapesAndTypes(const TypeProto& inferredType, TypeProto* existingType)
 /// ModelLocalFunctionsMap is a map of function id -> model local function proto
 /// All the ONNX helper utilities expect the function id == <function_proto.domain>:<function_proto.name>
 ///
-void InferShapes(
+ONNX_API void InferShapes(
     GraphProto* g,
     const std::unordered_map<std::string, int>& opset_imports,
     const ISchemaRegistry* schema_registry = OpSchemaRegistry::Instance(),
-    const ShapeInferenceOptions& options = {},
+    const ShapeInferenceOptions& options = ShapeInferenceOptions(),
     const ModelLocalFunctionsMap& in_model_functions = {});
 
-void InferShapes(
+ONNX_API void InferShapes(
     ModelProto& m,
     const ISchemaRegistry* schema_registry = OpSchemaRegistry::Instance(),
-    const ShapeInferenceOptions& options = {},
+    const ShapeInferenceOptions& options = ShapeInferenceOptions(),
     DataValueMap* generated_shape_data_by_name = nullptr);
 
-void InferShapes(
+ONNX_API void InferShapes(
     const std::string& model_path,
     const std::string& save_path = "",
     const ISchemaRegistry* schema_registry = OpSchemaRegistry::Instance(),
-    const ShapeInferenceOptions& options = {},
+    const ShapeInferenceOptions& options = ShapeInferenceOptions(),
     DataValueMap* generated_shape_data_by_name = nullptr);
 
 ///
@@ -473,7 +505,7 @@ void InferShapeForFunctionNode(
     const FunctionProto& func,
     const ISchemaRegistry* schema_registry,
     InferenceContext& ctx,
-    const ShapeInferenceOptions& options = {},
+    const ShapeInferenceOptions& options = ShapeInferenceOptions(),
     const ModelLocalFunctionsMap& model_local_functions_map = {},
     SymbolTable* symbolTable = nullptr,
     DataValueMap* generated_shape_data_by_name = nullptr);
@@ -487,7 +519,7 @@ void InferShapeForFunctionNode(
     const std::unordered_map<std::string, int>& func_opset_imports,
     const ISchemaRegistry* schema_registry,
     InferenceContext& ctx,
-    const ShapeInferenceOptions& options = {},
+    const ShapeInferenceOptions& options = ShapeInferenceOptions(),
     const ModelLocalFunctionsMap& model_local_functions_map = {},
     SymbolTable* symbolTable = nullptr,
     DataValueMap* generated_shape_data_by_name = nullptr);

onnx/shape_inference.py

@@ -10,12 +10,15 @@ complete.
 from __future__ import annotations
 
 import os
-from typing import Sequence
+from typing import TYPE_CHECKING
 
 import onnx
 import onnx.onnx_cpp2py_export.shape_inference as C  # noqa: N812
 from onnx import AttributeProto, FunctionProto, ModelProto, TypeProto
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 def infer_shapes(
     model: ModelProto | bytes,
@@ -105,7 +108,7 @@ def infer_node_outputs(
     opset_imports: list[onnx.OperatorSetIdProto] | None = None,
     ir_version: int = onnx.IR_VERSION,
 ) -> dict[str, onnx.TypeProto]:
-    if not schema.has_type_and_shape_inference_function:  # type: ignore
+    if not schema.has_type_and_shape_inference_function:
         return {}
     if input_data is None:
         input_data = {}
@@ -118,12 +121,12 @@ def infer_node_outputs(
 
     # catch KeyError if node's input does not exist in input_types
     passed_input_types = {
-        key: input_types[key].SerializeToString() for key in node.input
+        key: input_types[key].SerializeToString() for key in node.input if key != ""
     }
     # input_types will also be used as outer_scope_value_types so do not filter by node's input here
-    for key in input_types:
+    for key, value in input_types.items():
         if key not in passed_input_types:
-            passed_input_types[key] = input_types[key].SerializeToString()
+            passed_input_types[key] = value.SerializeToString()
     passed_input_data = {
         key: input_data[key].SerializeToString()
         for key in node.input