onnx 1.17.0__cp39-cp39-macosx_12_0_universal2.whl → 1.19.0__cp39-cp39-macosx_12_0_universal2.whl

onnx/backend/test/case/node/topk.py

@@ -10,8 +10,11 @@ from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 
 
-def topk_sorted_implementation(X, k, axis, largest):  # type: ignore
-    sorted_indices = np.argsort(X, axis=axis)
+def topk_sorted_implementation(X, k, axis, largest):
+    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.sort(X, axis=axis)
     if largest:
         sorted_indices = np.flip(sorted_indices, axis=axis)
@@ -55,6 +58,130 @@ class TopK(Base):
             node, inputs=[X, K], outputs=[values_ref, indices_ref], name="test_top_k"
         )
 
+    @staticmethod
+    def export_top_k_uint64() -> None:
+        axis = 1
+        largest = 1
+
+        k = 3
+        node = onnx.helper.make_node(
+            "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
+        )
+        X = np.array(
+            [
+                [0, 1, 2, 3],
+                [4, 5, 6, 7],
+                [8, 9, 10, 11],
+            ],
+            dtype=np.uint64,
+        )
+        K = np.array([k], dtype=np.int64)
+        values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)
+
+        # print(values_ref)
+        # [[ 3  2  1]
+        # [ 7  6  5]
+        # [11 10  9]]
+        # print(indices_ref)
+        # [[3 2 1]
+        # [3 2 1]
+        # [3 2 1]]
+
+        expect(
+            node,
+            inputs=[X, K],
+            outputs=[values_ref, indices_ref],
+            name="test_top_k_uint64",
+        )
+
+    @staticmethod
+    def export_top_k_same_values() -> None:
+        axis = 0
+        largest = 0
+
+        k = 3
+        node = onnx.helper.make_node(
+            "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
+        )
+        X = np.array(
+            [0, 0, 0, 0],
+            dtype=np.int64,
+        )
+        K = np.array([k], dtype=np.int64)
+        values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)
+
+        # (Pdb) print(values_ref)
+        # [0 0 0]
+        # (Pdb) print(indices_ref)
+        # [0 1 2]
+
+        expect(
+            node,
+            inputs=[X, K],
+            outputs=[values_ref, indices_ref],
+            name="test_top_k_same_values",
+        )
+
+    @staticmethod
+    def export_top_k_same_values_largest() -> None:
+        axis = 0
+        largest = 1
+
+        k = 3
+        node = onnx.helper.make_node(
+            "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
+        )
+        X = np.array(
+            [0, 0, 0, 0],
+            dtype=np.int64,
+        )
+        K = np.array([k], dtype=np.int64)
+        values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)
+
+        # print(values_ref)
+        # [0 0 0]
+        # print(indices_ref)
+        # [0 1 2]
+
+        expect(
+            node,
+            inputs=[X, K],
+            outputs=[values_ref, indices_ref],
+            name="test_top_k_same_values_largest",
+        )
+
+    @staticmethod
+    def export_top_k_same_values_2d() -> None:
+        axis = 1
+        largest = 1
+
+        k = 3
+        node = onnx.helper.make_node(
+            "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
+        )
+        X = np.array(
+            [[0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 1, 1]],
+            dtype=np.int64,
+        )
+        K = np.array([k], dtype=np.int64)
+        values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)
+
+        # print(values_ref)
+        # [[0 0 0]
+        # [1 1 1]
+        # [1 1 2]]
+        # print(indices_ref)
+        # [[0 1 2]
+        # [0 1 2]
+        # [2 3 0]]
+
+        expect(
+            node,
+            inputs=[X, K],
+            outputs=[values_ref, indices_ref],
+            name="test_top_k_same_values_2d",
+        )
+
     @staticmethod
     def export_top_k_smallest() -> None:
         axis = 1

onnx/backend/test/case/node/trilu.py

@@ -10,11 +10,11 @@ from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 
 
-def triu_reference_implementation(x, k=0):  # type: ignore
+def triu_reference_implementation(x, k=0):
     return np.triu(x, k)
 
 
-def tril_reference_implementation(x, k=0):  # type: ignore
+def tril_reference_implementation(x, k=0):
     return np.tril(x, k)
 
 

onnx/backend/test/case/node/unique.py

@@ -10,7 +10,7 @@ from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 
 
-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),
@@ -100,7 +100,7 @@ class Unique(Base):
             indices, inverse_indices, counts
         )
         # behavior changed with numpy >= 2.0
-        inverse_indices = inverse_indices.squeeze()
+        inverse_indices = inverse_indices.reshape(-1)
         # print(y)
         # [[1. 0. 0.]
         #  [2. 3. 4.]]
@@ -140,7 +140,7 @@ class Unique(Base):
             indices, inverse_indices, counts
         )
         # behavior changed with numpy >= 2.0
-        inverse_indices = inverse_indices.squeeze()
+        inverse_indices = inverse_indices.reshape(-1)
         # print(y)
         # [[[0. 1.]
         #  [1. 1.]
@@ -177,7 +177,7 @@ class Unique(Base):
             indices, inverse_indices, counts
         )
         # behavior changed with numpy >= 2.0
-        inverse_indices = inverse_indices.squeeze()
+        inverse_indices = inverse_indices.reshape(-1)
         # print(y)
         # [[0. 1.]
         #  [0. 1.]
@@ -195,3 +195,35 @@ class Unique(Base):
             outputs=[y, indices, inverse_indices, counts],
             name="test_unique_sorted_with_negative_axis",
         )
+
+    @staticmethod
+    def export_length_1() -> None:
+        node_sorted = onnx.helper.make_node(
+            "Unique",
+            inputs=["X"],
+            outputs=["Y", "indices", "inverse_indices", "counts"],
+            sorted=1,
+        )
+
+        x = np.array([0], dtype=np.int64)
+        y, indices, inverse_indices, counts = np.unique(x, True, True, True)
+        indices, inverse_indices, counts = specify_int64(
+            indices, inverse_indices, counts
+        )
+        # behavior changed with numpy >= 2.0
+        inverse_indices = inverse_indices.reshape(-1)
+        # print(y)
+        # [0]
+        # print(indices)
+        # [0]
+        # print(inverse_indices)
+        # [0]
+        # print(counts)
+        # [1]
+
+        expect(
+            node_sorted,
+            inputs=[x],
+            outputs=[y, indices, inverse_indices, counts],
+            name="test_unique_length_1",
+        )

onnx/backend/test/case/test_case.py

@@ -4,11 +4,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Sequence
+from typing import TYPE_CHECKING
 
-import numpy as np
+if TYPE_CHECKING:
+    from collections.abc import Sequence
 
-import onnx
+    import numpy as np
+
+    import onnx
 
 
 @dataclass

onnx/backend/test/case/utils.py

@@ -5,12 +5,15 @@ from __future__ import annotations
 
 import importlib
 import pkgutil
-from types import ModuleType
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from onnx import ONNX_ML
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
 all_numeric_dtypes = [
     np.int8,
     np.int16,
@@ -28,8 +31,7 @@ all_numeric_dtypes = [
 
 def import_recursive(package: ModuleType) -> None:
     """Takes a package and imports all modules underneath it."""
-    pkg_dir: list[str] | None = None
-    pkg_dir = package.__path__  # type: ignore
+    pkg_dir = package.__path__
     module_location = package.__name__
     for _module_loader, name, ispkg in pkgutil.iter_modules(pkg_dir):
         module_name = f"{module_location}.{name}"  # Module/package

onnx/backend/test/data/node/test_add_int8/test_data_set_0/output_0.pb

@@ -0,0 +1 @@
+BsumJ<+

onnx/backend/test/data/node/test_add_uint8/test_data_set_0/output_0.pb

@@ -1 +1,3 @@
-BsumJ<&,$!"' ,$'"
+BsumJ<*+%		#),#
+"#"#
+

onnx/backend/test/data/node/test_attention_3d_attn_mask/test_data_set_0/input_3.pb

@@ -0,0 +1 @@
+
B	attn_maskJ`pUw?��J?��,?O�z>��]>q*>�9l?7��>���>��>A6G?�X?i>5��>W�W?�jQ?N��=�" >>�U�=nm�>�f�=�v?ly|>

onnx/backend/test/data/node/test_attention_3d_attn_mask_expanded/test_data_set_0/input_3.pb

@@ -0,0 +1 @@
+
B	attn_maskJ`pUw?��J?��,?O�z>��]>q*>�9l?7��>���>��>A6G?�X?i>5��>W�W?�jQ?N��=�" >>�U�=nm�>�f�=�v?ly|>