onnx 1.17.0__cp312-cp312-macosx_12_0_universal2.whl → 1.19.0__cp312-cp312-macosx_12_0_universal2.whl

onnx/reference/ops/op_tensor_scatter.py

@@ -0,0 +1,50 @@
+# Copyright (c) ONNX Project Contributors
+
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import numpy as np
+
+from onnx.reference.op_run import OpRun
+
+
+class TensorScatter(OpRun):
+    def _run(self, past_cache, update, write_indices=None, mode="linear", axis=-2):
+        if mode not in {"linear", "circular"}:
+            raise ValueError(f"Unsupported mode: {mode}")
+
+        if write_indices is None:
+            write_indices = np.zeros((past_cache.shape[0],), dtype=np.int64)
+
+        input_shape = past_cache.shape
+        update_shape = update.shape
+        axis = axis % len(input_shape)
+
+        for i in range(len(input_shape)):
+            if i != axis:
+                if input_shape[i] != update_shape[i]:
+                    raise ValueError(
+                        f"Input shape {input_shape} and update shape {update_shape} are not compatible in dimension {i}"
+                    )
+            if i == axis:
+                if input_shape[i] < update_shape[i]:
+                    raise ValueError(
+                        f"Input shape {input_shape} and update shape {update_shape} are not compatible in axis dimension"
+                    )
+
+        max_sequence_length = input_shape[axis]
+        sequence_length = update_shape[axis]
+        present_cache = np.copy(past_cache)
+
+        for prefix_idx in np.ndindex(input_shape[:axis]):
+            batch_idx = prefix_idx[0]
+            for sequence_idx in range(sequence_length):
+                cache_idx = (*prefix_idx, write_indices[batch_idx] + sequence_idx)
+                if mode == "circular":
+                    cache_idx = tuple(
+                        np.mod(np.asarray(cache_idx), max_sequence_length)
+                    )
+                update_idx = (*prefix_idx, sequence_idx)
+                present_cache[cache_idx] = update[update_idx]
+
+        return (present_cache.reshape(input_shape),)

onnx/reference/ops/op_tfidf_vectorizer.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import collections
 from enum import IntEnum
 
 import numpy as np
@@ -10,9 +11,9 @@ import numpy as np
 from onnx.reference.op_run import OpRun
 
 
-class IntMap(dict):  # type: ignore
+class IntMap(collections.UserDict):
     def __init__(self):
-        dict.__init__(self)
+        super().__init__()
         self.added_keys = []
 
     def emplace(self, key, value):
@@ -22,8 +23,8 @@ class IntMap(dict):  # type: ignore
             raise TypeError(f"value must be a NGramPart not {type(value)}.")
         if key not in self:
             self.added_keys.append(key)
-            self[key] = value
-        return self[key]
+            self.data[key] = value
+        return self.data[key]
 
     def __repr__(self):
         vals = {k: repr(v) for k, v in self.items()}
@@ -59,7 +60,7 @@ class NgramPart:
         self._leafs_ = None
 
     def init(self):
-        self._leafs_ = IntMap()  # type: ignore
+        self._leafs_ = IntMap()
 
     def __repr__(self):
         if self.empty():
@@ -81,7 +82,7 @@ class NgramPart:
     def find(self, key):
         if not self.has_leaves():
             return None
-        if key in self._leafs_:  # type: ignore
+        if key in self._leafs_:
             return key
         return None
 
@@ -89,7 +90,7 @@ class NgramPart:
         return self.leafs_.emplace(key, value)
 
     def __getitem__(self, key):
-        return self._leafs_[key]  # type: ignore
+        return self._leafs_[key]
 
 
 class WeightingCriteria(IntEnum):
@@ -126,9 +127,9 @@ def populate_grams(
 
 
 class TfIdfVectorizer(OpRun):
-    def __init__(self, onnx_node, run_params):  # type: ignore
+    def __init__(self, onnx_node, run_params):
         OpRun.__init__(self, onnx_node, run_params)
-        mode = self.mode  # type: ignore
+        mode = self.mode
 
         if mode == "TF":
             self.weighting_criteria_ = WeightingCriteria.TF
@@ -137,16 +138,16 @@ class TfIdfVectorizer(OpRun):
         elif mode == "TFIDF":
             self.weighting_criteria_ = WeightingCriteria.TFIDF
 
-        self.min_gram_length_ = self.min_gram_length  # type: ignore
-        self.max_gram_length_ = self.max_gram_length  # type: ignore
-        self.max_skip_count_ = self.max_skip_count  # type: ignore
-        self.ngram_counts_ = self.ngram_counts  # type: ignore
-        self.max_gram_length_ = self.max_gram_length  # type: ignore
-        self.ngram_indexes_ = self.ngram_indexes  # type: ignore
+        self.min_gram_length_ = self.min_gram_length
+        self.max_gram_length_ = self.max_gram_length
+        self.max_skip_count_ = self.max_skip_count
+        self.ngram_counts_ = self.ngram_counts
+        self.max_gram_length_ = self.max_gram_length
+        self.ngram_indexes_ = self.ngram_indexes
         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.weights_ = self.weights
+        self.pool_int64s_ = self.pool_int64s
+        self.pool_strings_ = self.pool_strings
 
         self.int64_map_ = NgramPart(-10)
         self.int64_map_.init()
@@ -237,7 +238,7 @@ class TfIdfVectorizer(OpRun):
             raise RuntimeError("Unexpected weighting_criteria.")
         return Y.reshape(output_dims)
 
-    def compute_impl(  # type: ignore
+    def compute_impl(
         self,
         X: np.ndarray,
         row_num: int,
@@ -300,7 +301,7 @@ class TfIdfVectorizer(OpRun):
                 if start_ngram_size > max_gram_length:
                     break
 
-    def _run(  # type: ignore
+    def _run(
         self,
         X,
         max_gram_length=None,

onnx/reference/ops/op_thresholded_relu.py

@@ -9,6 +9,6 @@ from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class ThresholdedRelu(OpRunUnaryNum):
-    def _run(self, x, alpha=None):  # type: ignore
-        alpha = alpha or self.alpha  # type: ignore
-        return (np.where(x > alpha, x, 0).astype(x.dtype),)  # type: ignore
+    def _run(self, x, alpha=None):
+        alpha = alpha or self.alpha
+        return (np.where(x > alpha, x, 0).astype(x.dtype),)

onnx/reference/ops/op_tile.py

@@ -9,5 +9,5 @@ from onnx.reference.op_run import OpRun
 
 
 class Tile(OpRun):
-    def _run(self, x, repeats):  # type: ignore
+    def _run(self, x, repeats):
         return (np.tile(x, repeats),)

onnx/reference/ops/op_topk.py

@@ -8,7 +8,7 @@ import numpy as np
 from onnx.reference.op_run import OpRun
 
 
-def topk_sorted_implementation(X, k, axis, largest):  # type: ignore
+def topk_sorted_implementation(X, k, axis, largest):
     """See function `_kneighbors_reduce_func
     <https://github.com/scikit-learn/scikit-learn/blob/main/
     sklearn/neighbors/_base.py#L304>`_.
@@ -19,27 +19,12 @@ def topk_sorted_implementation(X, k, axis, largest):  # type: ignore
         k = k[0]
     # This conversion is needed for distribution x86.
     k = int(k)
-    if len(X.shape) == 2 and axis == 1:
-        sample_range = np.arange(X.shape[0])[:, None]
-        if largest == 0:
-            sorted_indices = np.argpartition(X, axis=axis, kth=k - 1)
-            sorted_indices = sorted_indices[:, :k]
-            # argpartition doesn't guarantee sorted order, so we sort again
-            sorted_indices = sorted_indices[
-                sample_range, np.argsort(X[sample_range, sorted_indices])
-            ]
-        else:
-            sorted_indices = np.argpartition(-X, axis=axis, kth=k - 1)
-            sorted_indices = sorted_indices[:, :k]
-            # argpartition doesn't guarantee sorted order, so we sort again
-            sorted_indices = sorted_indices[
-                sample_range, np.argsort(-X[sample_range, sorted_indices])
-            ]
-        sorted_distances = X[sample_range, sorted_indices]
-        return sorted_distances, sorted_indices
-
-    sorted_indices = np.argsort(X, axis=axis)
-    sorted_values = np.sort(X, axis=axis)
+    # Used to tiebreak
+    ind_axis = np.indices(X.shape)[axis]
+    if largest:
+        ind_axis = -ind_axis
+    sorted_indices = np.lexsort((ind_axis, X), axis=axis)
+    sorted_values = np.take_along_axis(X, sorted_indices, axis=axis)
     if largest:
         sorted_indices = np.flip(sorted_indices, axis=axis)
         sorted_values = np.flip(sorted_values, axis=axis)
@@ -50,7 +35,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
+    def _common_run(self, data, ink, axis, largest=1):
         """Runtime for operator *TopK*.
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
@@ -63,13 +48,13 @@ class _CommonTopK(OpRun):
             <https://github.com/Microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cpu/math/top_k.cc#L63>`_.
         """
         k = ink[0]
-        axis = axis if axis >= 0 else (axis + len(data.shape))  # type: ignore
+        axis = axis if axis >= 0 else (axis + len(data.shape))
         sort, sorti = topk_sorted_implementation(data, k, axis, largest)
         return (sort, sorti.astype(np.int64))
 
 
 class TopK_1(_CommonTopK):
-    def _run(self, data, k=None, axis=None):  # type: ignore
+    def _run(self, data, k=None, axis=None):
         """Runtime for operator *TopK*.
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
@@ -81,11 +66,11 @@ class TopK_1(_CommonTopK):
             does in `top_k.cc
             <https://github.com/Microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cpu/math/top_k.cc#L63>`_.
         """
-        return _CommonTopK._common_run(self, data, [k], axis=axis)  # type: ignore
+        return _CommonTopK._common_run(self, data, [k], axis=axis)
 
 
 class TopK_10(_CommonTopK):
-    def _run(self, data, ink, axis=None):  # type: ignore
+    def _run(self, data, ink, axis=None):
         """Runtime for operator *TopK*.
         The implementation is not the most efficient
         as it sorts everything then extracts the top *k*
@@ -108,7 +93,7 @@ class TopK_11(_CommonTopK):
         axis=None,
         largest=None,
         sorted=None,  # noqa: A002
-    ):  # type: ignore
+    ):
         """Runtime for operator *TopK*.
 
         The implementation is not the most efficient

onnx/reference/ops/op_transpose.py

@@ -9,7 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 class Transpose(OpRun):
-    def _run(self, data, perm=None):  # type: ignore
+    def _run(self, data, perm=None):
         perm_ = None if (perm is None or len(perm) == 0) else perm
         if perm_ is None:
             return (np.transpose(data),)

onnx/reference/ops/op_trilu.py

@@ -9,8 +9,8 @@ from onnx.reference.op_run import OpRun
 
 
 class Trilu(OpRun):
-    def _run(self, x, k=None, upper=None):  # type: ignore
-        k = 0 if k is None else int(k)
-        if upper:  # type: ignore
+    def _run(self, x, k=None, upper=None):
+        k = 0 if k is None else k.item()
+        if upper:
             return (np.triu(x, k),)
         return (np.tril(x, k).astype(x.dtype),)

onnx/reference/ops/op_unique.py

@@ -8,7 +8,7 @@ import numpy as np
 from onnx.reference.op_run import OpRun
 
 
-def _specify_int64(indices, inverse_indices, counts):  # type: ignore
+def _specify_int64(indices, inverse_indices, counts):
     return (
         np.array(indices, dtype=np.int64),
         np.array(inverse_indices, dtype=np.int64),
@@ -43,7 +43,7 @@ class Unique(OpRun):
             indices, inverse_indices, counts
         )
         # numpy 2.0 has a different behavior than numpy 1.x.
-        inverse_indices = inverse_indices.squeeze()
+        inverse_indices = inverse_indices.reshape(-1)
         if len(self.onnx_node.output) == 2:
             return (y, indices)
         if len(self.onnx_node.output) == 3:

onnx/reference/ops/op_unsqueeze.py

@@ -9,7 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 class Unsqueeze_1(OpRun):
-    def _run(self, data, axes=None):  # type: ignore
+    def _run(self, data, axes=None):
         if isinstance(axes, np.ndarray):
             axes = tuple(axes)
         elif axes in ([], ()):
@@ -32,7 +32,7 @@ class Unsqueeze_11(Unsqueeze_1):
 
 
 class Unsqueeze_13(OpRun):
-    def _run(self, data, axes=None):  # type: ignore
+    def _run(self, data, axes=None):
         if axes is not None:
             if hasattr(axes, "__iter__") and len(axes.shape) > 0:
                 try:

onnx/reference/ops/op_upsample.py

@@ -9,7 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 class Upsample(OpRun):
-    def _run(self, x, scale, mode=None):  # type: ignore
+    def _run(self, x, scale, mode=None):
         if mode == "nearest" and scale.astype(np.int64).tolist() == scale.tolist():
             r = x
             for axis, s in enumerate(scale):

onnx/reference/ops/op_where.py

@@ -9,7 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 class Where(OpRun):
-    def _run(self, condition, x, y):  # type: ignore
+    def _run(self, condition, x, y):
         if (
             x.dtype != y.dtype
             and x.dtype not in (object,)

onnx/reference/ops/op_xor.py

@@ -9,5 +9,5 @@ from onnx.reference.ops._op import OpRunBinary
 
 
 class Xor(OpRunBinary):
-    def _run(self, x, y):  # type: ignore
+    def _run(self, x, y):
         return (np.logical_xor(x, y),)

onnx/reference/ops_optimized/__init__.py

@@ -1,5 +1,3 @@
-# Copyright (c) ONNX Project Contributors
-
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0

onnx/reference/ops_optimized/op_conv_optimized.py

@@ -47,7 +47,7 @@ def im2col_fast(X, kernel_shape, pads, strides):
     return conc_cols, tuple(shape_out)
 
 
-def _conv_implementation_im2col(  # type: ignore
+def _conv_implementation_im2col(
     X, W, B, auto_pad, dilations, group, kernel_shape, pads, strides
 ):
     if dilations is None:
@@ -165,7 +165,7 @@ def _conv_implementation_im2col(  # type: ignore
 
 
 class Conv(OpRun):
-    def _run(  # type: ignore
+    def _run(
         self,
         X,
         W,