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

onnx/reference/ops/op_scatter_elements.py

@@ -9,13 +9,14 @@ from onnx.reference.op_run import OpRun
 
 
 def scatter_elements(data, indices, updates, axis=0, reduction=None):  # type: ignore
-    """
+    """Scatter elements.
+
     ::
-        // for 3-dim and axis=0
-        //    output[indices[i][j][k]][j][k] = updates[i][j][k]
-        // for axis 1
-        //    output[i][indices[i][j][k]][k] = updates[i][j][k]
-        // and so on
+        for 3-dim and axis=0
+            output[indices[i][j][k]][j][k] = updates[i][j][k]
+        for axis 1
+            output[i][indices[i][j][k]][k] = updates[i][j][k]
+        and so on.
     """
     if reduction == "add":
 

onnx/reference/ops/op_stft.py

@@ -28,8 +28,8 @@ def _unsqueeze(a, axis):  # type: ignore
 
 
 def _stft(x, fft_length: int, hop_length, n_frames, window, onesided=False):  # type: ignore
-    """
-    Applies one dimensional FFT with window weights.
+    """Applies one dimensional FFT with window weights.
+
     torch defines the number of frames as:
     `n_frames = 1 + (len - n_fft) // hop_length`.
     """
@@ -75,9 +75,7 @@ def _stft(x, fft_length: int, hop_length, n_frames, window, onesided=False):  #
 
 
 def _istft(x, fft_length: int, hop_length, window, onesided=False):  # type: ignore
-    """
-    Reverses of `stft`.
-    """
+    """Reverses of `stft`."""
     zero = [0]
     one = [1]
     two = [2]

onnx/reference/ops/op_string_normalizer.py

@@ -13,8 +13,7 @@ from onnx.reference.op_run import OpRun, RuntimeTypeError
 
 
 class StringNormalizer(OpRun):
-    """
-    The operator is not really threadsafe as python cannot
+    """The operator is not really threadsafe as python cannot
     play with two locales at the same time. stop words
     should not be implemented here as the tokenization
     usually happens after this steps.
@@ -129,15 +128,16 @@ class StringNormalizer(OpRun):
 
     @staticmethod
     def strip_accents_unicode(s):  # type: ignore
-        """
-        Transforms accentuated unicode symbols into their simple counterpart.
+        """Transforms accentuated unicode symbols into their simple counterpart.
         Source: `sklearn/feature_extraction/text.py
         <https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/
         feature_extraction/text.py#L115>`_.
 
-        :param s: string
-            The string to strip
-        :return: the cleaned string
+        Args:
+            s: string The string to strip
+
+        Returns:
+            the cleaned string
         """
         try:
             # If `s` is ASCII-compatible, then it does not contain any accented

onnx/reference/ops/op_tfidf_vectorizer.py

@@ -17,8 +17,8 @@ class IntMap(dict):  # type: ignore
         self.added_keys = []
 
     def emplace(self, key, value):
-        if not isinstance(key, int):
-            raise TypeError(f"key must be a NGramPart not {type(key)}.")
+        if not isinstance(key, (int, str)):
+            raise TypeError(f"key must be a int or str not {type(key)}.")
         if not isinstance(value, NgramPart):
             raise TypeError(f"value must be a NGramPart not {type(value)}.")
         if key not in self:
@@ -147,11 +147,12 @@ class TfIdfVectorizer(OpRun):
         self.output_size_ = max(self.ngram_indexes_) + 1
         self.weights_ = self.weights  # type: ignore
         self.pool_int64s_ = self.pool_int64s  # type: ignore
+        self.pool_strings_ = self.pool_strings  # type: ignore
 
         self.int64_map_ = NgramPart(-10)
         self.int64_map_.init()
 
-        total_items = len(self.pool_int64s_)
+        total_items = len(self.pool_int64s_ or self.pool_strings_)
         ngram_id = 1  # start with 1, 0 - means no n-gram
         # Load into dictionary only required gram sizes
         ngram_size = 1
@@ -170,7 +171,7 @@ class TfIdfVectorizer(OpRun):
                     and ngram_size <= self.max_gram_length_
                 ):
                     ngram_id = populate_grams(
-                        self.pool_int64s_,
+                        self.pool_int64s_ or self.pool_strings_,
                         start_idx,
                         ngrams,
                         ngram_size,
@@ -207,10 +208,8 @@ class TfIdfVectorizer(OpRun):
 
         w = self.weights_
         if self.weighting_criteria_ == WeightingCriteria.TF:
-            i = 0
-            for f in frequencies:
+            for i, f in enumerate(frequencies):
                 Y[i] = f
-                i += 1
         elif self.weighting_criteria_ == WeightingCriteria.IDF:
             if len(w) > 0:
                 p = 0
@@ -359,7 +358,7 @@ class TfIdfVectorizer(OpRun):
             # TfidfVectorizer returns a zero tensor of shape
             # {b_dim, output_size} when b_dim is the number of received observations
             # and output_size the is the maximum value in ngram_indexes attribute plus 1.
-            return self.output_result(B, frequencies)  # type: ignore[arg-type]
+            return (self.output_result(B, frequencies),)  # type: ignore[arg-type]
 
         def fn(row_num):
             self.compute_impl(

onnx/reference/ops/op_topk.py

@@ -9,8 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 def topk_sorted_implementation(X, k, axis, largest):  # type: ignore
-    """
-    See function `_kneighbors_reduce_func
+    """See function `_kneighbors_reduce_func
     <https://github.com/scikit-learn/scikit-learn/blob/main/
     sklearn/neighbors/_base.py#L304>`_.
     """
@@ -52,8 +51,7 @@ def topk_sorted_implementation(X, k, axis, largest):  # type: ignore
 
 class _CommonTopK(OpRun):
     def _common_run(self, data, ink, axis, largest=1):  # type: ignore
-        """
-        Runtime for operator *TopK*.
+        """Runtime for operator *TopK*.
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
         values.
@@ -72,8 +70,7 @@ class _CommonTopK(OpRun):
 
 class TopK_1(_CommonTopK):
     def _run(self, data, k=None, axis=None):  # type: ignore
-        """
-        Runtime for operator *TopK*.
+        """Runtime for operator *TopK*.
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
         values.
@@ -89,8 +86,7 @@ class TopK_1(_CommonTopK):
 
 class TopK_10(_CommonTopK):
     def _run(self, data, ink, axis=None):  # type: ignore
-        """
-        Runtime for operator *TopK*.
+        """Runtime for operator *TopK*.
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
         values.
@@ -105,9 +101,11 @@ class TopK_10(_CommonTopK):
 
 
 class TopK_11(_CommonTopK):
-    def _run(self, data, ink, axis=None, largest=None, sorted=None):  # type: ignore
-        """
-        Runtime for operator *TopK*.
+    def _run(
+        self, data, ink, axis=None, largest=None, sorted=None  # noqa: A002
+    ):  # type: ignore
+        """Runtime for operator *TopK*.
+
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
         values.

onnx/reference/ops/op_unique.py

@@ -17,7 +17,7 @@ def _specify_int64(indices, inverse_indices, counts):  # type: ignore
 
 
 class Unique(OpRun):
-    def _run(self, x, axis=None, sorted=None):  # type: ignore
+    def _run(self, x, axis=None, sorted=None):  # type: ignore  # noqa: A002
         if axis is None or np.isnan(axis):
             y, indices, inverse_indices, counts = np.unique(x, True, True, True)
         else:

onnx/reference/reference_evaluator.py

@@ -9,7 +9,16 @@ import numpy as np
 
 from onnx import load
 from onnx.defs import onnx_opset_version
-from onnx.onnx_pb import FunctionProto, GraphProto, ModelProto, NodeProto, TypeProto
+from onnx.external_data_helper import ExternalDataInfo, uses_external_data
+from onnx.model_container import ModelContainer
+from onnx.onnx_pb import (
+    FunctionProto,
+    GraphProto,
+    ModelProto,
+    NodeProto,
+    TensorProto,
+    TypeProto,
+)
 from onnx.reference.op_run import (
     OpFunctionContextDependant,
     OpRun,
@@ -21,36 +30,36 @@ from onnx.reference.ops_optimized import optimized_operators
 
 
 class ReferenceEvaluator:
-    """
-    Computes the outputs of an ONNX proto
-    (`ModelProto`, `FunctionProto`, `GraphProto`, `NodeProto`).
+    r"""Computes the outputs of an ONNX proto (`ModelProto`, `FunctionProto`, `GraphProto`, `NodeProto`).
+
     This is a pure python implementation of ONNX specifications.
     Mismatches may remain between the official specifications and the implementation here.
     In the case of such a mismatch, the official spec overrides this implementation.
 
-    :param proto: :class:`onnx.ModelProto`, :class:`onnx.GraphProto`,
-        :class:`onnx.FunctionProto`, :class:`onnx.NodeProto`,
-        filename or bytes
-    :param verbose: display intermediate results
-        on the standard output during the execution
-    :param opsets: if *proto* is an instance of *GraphProto*,
-        opsets must be defined by a dictionary of
-    :param functions: known onnx functions
-    :param new_ops: this runtime can be used to test the implementations
-        of new operators, *new_ops* is a list of classes
-        derived from :class:`OpRun <onnx.reference.op_run.OpRun>`,
-        every class must define the static attribute `domain`,
-        there may be multiple implementations for the same operator,
-        the first one in the list is used.
-    :param optimized: some operators have two implementations,
-        a naive one corresponding to definition of the mathematical
-        definition of the operator, another one more efficient.
-        This is the case for operator Conv. The naive version is ten times
-        slower than the optimized one using a decomposition
-        into *Conv = im2col + Gemm*. If True, all optimized
-        kernels are added in `new_ops` and are used instead of the
-        inner implementation if list *new_ops* does not already contain
-        one.
+    Args:
+        proto: :class:`onnx.ModelProto`, :class:`onnx.GraphProto`,
+            :class:`onnx.FunctionProto`, :class:`onnx.NodeProto`,
+            filename or bytes
+        verbose: display intermediate results on the standard output
+            during the execution
+        opsets: if *proto* is an instance of *GraphProto*, opsets must
+            be defined by a dictionary of
+        functions: known onnx functions
+        new_ops: this runtime can be used to test the implementations of
+            new operators, *new_ops* is a list of classes derived from
+            :class:`OpRun <onnx.reference.op_run.OpRun>`, every class
+            must define the static attribute `domain`, there may be
+            multiple implementations for the same operator, the first
+            one in the list is used.
+        optimized: some operators have two implementations, a naive one
+            corresponding to definition of the mathematical definition
+            of the operator, another one more efficient. This is the
+            case for operator Conv. The naive version is ten times
+            slower than the optimized one using a decomposition into
+            *Conv = im2col + Gemm*. If True, all optimized kernels are
+            added in `new_ops` and are used instead of the inner
+            implementation if list *new_ops* does not already contain
+            one.
 
     The class maps every node to its associated implementation.
     When a subgraph of a function is met,
@@ -209,6 +218,13 @@ class ReferenceEvaluator:
                         new_ops.append(op)
         self.output_types_ = None
         self.input_types_ = None
+
+        if isinstance(proto, ModelContainer):
+            self.container_ = proto
+            proto = self.container_.model_proto
+        else:
+            self.container_ = None
+
         if isinstance(proto, str):
             with open(proto, "rb") as f:
                 proto = load(f)
@@ -268,9 +284,8 @@ class ReferenceEvaluator:
         if functions is not None:
             for f in functions:  # type: ignore
                 if isinstance(f, FunctionProto):
-                    existing_functions = list(self.functions_.values())
-                    self.functions_[f.domain, f.name] = ReferenceEvaluator(
-                        f, verbose=verbose, functions=existing_functions
+                    self.functions_[f.domain, f.name] = self.__class__(
+                        f, verbose=verbose, functions=list(self.functions_.values())
                     )
                 elif isinstance(f, ReferenceEvaluator):
                     onx = f.proto_  # type: ignore
@@ -294,6 +309,24 @@ class ReferenceEvaluator:
                 self.new_ops_[key] = cl
         self._init()
 
+    def retrieve_external_data(self, initializer: TensorProto) -> np.array:
+        """Returns a tensor saved as external."""
+        info = ExternalDataInfo(initializer)
+        location = info.location
+        if self.container_ and self.container_.is_in_memory_external_initializer(
+            location
+        ):
+            # It comes from a large container.
+            return self.container_[location]
+        # Otherwise, the data is on disk.
+        if self.container_ is not None:
+            raise RuntimeError(
+                "ReferenceEvaluator assumes a LargeContainer was loaded with its external tensor."
+            )
+        raise RuntimeError(
+            "An instance of LargeContainer should be created before using ReferenceEvaluator."
+        )
+
     def _log_arg(self, a: Any) -> Any:
         if isinstance(a, (str, int, float)):
             return a
@@ -316,33 +349,32 @@ class ReferenceEvaluator:
 
     @property
     def input_names(self):  # type: ignore
-        "Returns the input names."
+        """Returns the input names."""
         return self.input_names_
 
     @property
     def input_types(self):  # type: ignore
-        "Returns the input types if any specified."
+        """Returns the input types if any specified."""
         return self.input_types_
 
     @property
     def output_names(self):  # type: ignore
-        "Returns the output names."
+        """Returns the output names."""
         return self.output_names_
 
     @property
     def output_types(self):  # type: ignore
-        "Returns the output types."
+        """Returns the output types."""
         return self.output_types_
 
     @property
     def opsets(self):  # type: ignore
-        "Returns the opsets."
+        """Returns the opsets."""
         return self.opsets_
 
     @property
     def has_linked_attribute(self):
-        """
-        Checks if the graph has a linked attribute (= an attribute whose value is defined
+        """Checks if the graph has a linked attribute (= an attribute whose value is defined
         by a function attribute.
         """
         return any(node.has_linked_attribute for node in self.rt_nodes_)
@@ -364,18 +396,22 @@ class ReferenceEvaluator:
         return self.all_types_[name]
 
     def _init(self) -> None:
-        """
-        Loads the implementation for every node in the graph.
-        """
+        """Loads the implementation for every node in the graph."""
         self.rt_inits_ = {}
         self.rt_nodes_ = []
         for init in self.inits_:
-            self.rt_inits_[init.name] = to_array_extended(init)  # type: ignore[union-attr,arg-type]
+            self.rt_inits_[init.name] = (
+                self.retrieve_external_data(init)
+                if uses_external_data(init)
+                else to_array_extended(init)
+            )
         run_params = {
             "log": lambda pattern, *args: self._log(10, pattern, *args),
             "opsets": self.opsets,
             "verbose": self.verbose,
             "new_ops": self.new_ops_,
+            "existing_functions": self.functions_.copy(),
+            "evaluator_cls": self.__class__,
         }
         if self.input_types_:
             all_types = {i.name: i.type for i in self.onnx_graph_.input}
@@ -419,9 +455,7 @@ class ReferenceEvaluator:
     def _load_impl(  # noqa: PLR0911
         self, node: NodeProto, input_types: Optional[TypeProto] = None
     ) -> Any:
-        """
-        Loads the implementation for a specified runtime.
-        """
+        """Loads the implementation for a specified runtime."""
         if node.domain not in self.opsets:
             raise RuntimeError(
                 f"Domain {node.domain!r} (node type: {node.op_type!r}) "
@@ -444,7 +478,13 @@ class ReferenceEvaluator:
             from onnx.reference.ops import load_op
 
             try:
-                return load_op(node.domain, node.op_type, version, expand=expand)
+                return load_op(
+                    node.domain,
+                    node.op_type,
+                    version,
+                    expand=expand,
+                    evaluator_cls=self.__class__,
+                )
             except RuntimeContextError:
                 if input_types is None:
                     raise
@@ -455,6 +495,7 @@ class ReferenceEvaluator:
                     node=node,
                     input_types=input_types,  # type: ignore[arg-type]
                     expand=expand,
+                    evaluator_cls=self.__class__,
                 )
 
         if expand:
@@ -466,38 +507,52 @@ 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)
+            return load_op_pt(
+                node.domain, node.op_type, version, evaluator_cls=self.__class__
+            )
 
         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)
+            return load_op_exp(
+                node.domain, node.op_type, version, evaluator_cls=self.__class__
+            )
 
         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)
+            return load_op_ml(
+                node.domain, node.op_type, version, evaluator_cls=self.__class__
+            )
 
         # It has to be a function.
         if key in self.functions_:
             from onnx.reference.ops import load_op
 
             impl = self.functions_[key]
-            return load_op(node.domain, node.op_type, version, custom=impl)
+            return load_op(
+                node.domain,
+                node.op_type,
+                version,
+                custom=impl,
+                evaluator_cls=self.__class__,
+            )
         raise NotImplementedError(
             f"Node type {node.op_type!r} from domain {node.domain!r} "
             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
-        """
-        Executes the onnx model.
+        """Executes the onnx model.
 
-        :param output_names: requested outputs by names,
-            None for all
-        :param feed_inputs: dictionary `{ input name: input value }`
-        :param attributes: attributes value if the instance runs a FunctionProto
-        :return: list of requested outputs
+        Args:
+            output_names: requested outputs by names, None for all
+            feed_inputs: dictionary `{ input name: input value }`
+            attributes: attributes value if the instance runs a
+                FunctionProto
+
+        Returns:
+            list of requested outputs
         """
         if output_names is None:
             output_names = self.output_names
@@ -516,6 +571,13 @@ class ReferenceEvaluator:
         # step 2: execute nodes
         for node in self.rt_nodes_:
             self._log(1, "%s(%s) -> %s", node.op_type, node.input, node.output)
+            for i in node.input:
+                if i not in results:
+                    raise RuntimeError(
+                        f"Unable to find input {i!r} in known results {sorted(results)}, "
+                        f"self.rt_inits_ has {sorted(self.rt_inits_)}, "
+                        f"feed_inputs has {sorted(feed_inputs)}."
+                    )
             inputs = [results[i] for i in node.input]
             linked_attributes = {}
             if node.has_linked_attribute and attributes:
@@ -525,19 +587,13 @@ class ReferenceEvaluator:
             else:
                 outputs = node.run(*inputs, **linked_attributes)
             for name, value in zip(node.output, outputs):
-                if isinstance(value, tuple):
-                    raise TypeError(
-                        f"Unexected type {type(value)} for output {name!r}."
-                    )
                 self._log(2, " + %s: %s", name, value)  # type: ignore[arg-type]
                 results[name] = value
 
         # return the results
-        list_results: List[Any] = []
         for name in output_names:
             if name not in results:
                 raise RuntimeError(
                     f"Unable to find output name {name!r} in {sorted(results)}, proto is\n{self.proto_}"
                 )
-            list_results.append(results[name])
-        return list_results
+        return [results[name] for name in output_names]