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

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

@@ -0,0 +1,231 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import numpy as np
+
+import onnx
+from onnx.backend.test.case.base import Base
+from onnx.backend.test.case.node import expect
+from onnx.reference.ops.op_rotary_embedding import rotary_embedding
+
+
+class RotaryEmbedding(Base):
+    @staticmethod
+    def export_rotary_embedding() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache", "position_ids"],
+            outputs=["output"],
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
+        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data, cos_cache_data, sin_cache_data, position_ids=position_ids_data
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_3d_input() -> None:
+        num_heads = 4
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache", "position_ids"],
+            outputs=["output"],
+            num_heads=num_heads,
+        )
+
+        input_data = np.random.rand(2, 3, 32).astype(np.float32)
+        position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
+        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data,
+            cos_cache_data,
+            sin_cache_data,
+            position_ids=position_ids_data,
+            num_heads=num_heads,
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_3d_input",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_interleaved() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache", "position_ids"],
+            outputs=["output"],
+            interleaved=1,
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
+        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data,
+            cos_cache_data,
+            sin_cache_data,
+            position_ids=position_ids_data,
+            interleaved=1,
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_interleaved",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_with_rotary_dim() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache", "position_ids"],
+            outputs=["output"],
+            rotary_embedding_dim=4,
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
+        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data,
+            cos_cache_data,
+            sin_cache_data,
+            position_ids=position_ids_data,
+            rotary_embedding_dim=4,
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_with_rotary_dim",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_with_interleaved_rotary_dim() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache", "position_ids"],
+            outputs=["output"],
+            rotary_embedding_dim=4,
+            interleaved=1,
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
+        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data,
+            cos_cache_data,
+            sin_cache_data,
+            position_ids=position_ids_data,
+            interleaved=1,
+            rotary_embedding_dim=4,
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_with_interleaved_rotary_dim",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_no_position_ids() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache"],
+            outputs=["output"],
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        sin_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(input_data, cos_cache_data, sin_cache_data)
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_no_position_ids",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_no_position_ids_interleaved() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache"],
+            outputs=["output"],
+            interleaved=1,
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        sin_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data,
+            cos_cache_data,
+            sin_cache_data,
+            interleaved=1,
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_no_position_ids_interleaved",
+        )
+
+    @staticmethod
+    def export_rotary_embedding_no_position_ids_rotary_dim() -> None:
+        node = onnx.helper.make_node(
+            "RotaryEmbedding",
+            inputs=["input", "cos_cache", "sin_cache"],
+            outputs=["output"],
+            rotary_embedding_dim=4,
+        )
+
+        input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
+        sin_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+        cos_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+
+        expected_output = rotary_embedding(
+            input_data,
+            cos_cache_data,
+            sin_cache_data,
+            rotary_embedding_dim=4,
+        )
+
+        expect(
+            node,
+            inputs=[input_data, cos_cache_data, sin_cache_data],
+            outputs=[expected_output],
+            name="test_rotary_embedding_no_position_ids_rotary_dim",
+        )

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

@@ -29,11 +29,36 @@ class Sub(Base):
         z = x - y
         expect(node, inputs=[x, y], outputs=[z], name="test_sub")
 
+        x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.int8)
+        y = np.random.randint(12, size=(3, 4, 5), dtype=np.int8)
+        z = x - y
+        expect(node, inputs=[x, y], outputs=[z], name="test_sub_int8")
+
+        x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.int16)
+        y = np.random.randint(12, size=(3, 4, 5), dtype=np.int16)
+        z = x - y
+        expect(node, inputs=[x, y], outputs=[z], name="test_sub_int16")
+
         x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint8)
         y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint8)
         z = x - y
         expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint8")
 
+        x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint16)
+        y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint16)
+        z = x - y
+        expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint16")
+
+        x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint32)
+        y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint32)
+        z = x - y
+        expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint32")
+
+        x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint64)
+        y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint64)
+        z = x - y
+        expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint64")
+
     @staticmethod
     def export_sub_broadcast() -> None:
         node = onnx.helper.make_node(

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_acosh/test_data_set_0/output_0.pb

@@ -1 +1 @@
-
B
yJ�
��@D{,@�#@��@�8@��&@w�@��8@[m=@@�22@*�@
+
B
yJ�
��@D{,@�#@��@�8@��&@w�@��8@Zm=@@�22@*�@

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

@@ -1 +1 @@
-
B
yJ�
ø?O
L?(�%?�?���>��3?J��>���?8s�?l{�>��i?�? �?�_�?���=ت�=�<��{?"Hd?��?���?fm?C��>E@e?[��=��1?M>`H�?T�?���>��>Z�b?%��>$�?c�<tm*?�(?3*?��?B@?{H�>]��>k�E?9�v=:�:?p-<?�
+
B
yJ�
ĸ?P
L?*�%?�?���>��3?J��>���?8s�?l{�>��i?�? �?�_�?���=ت�=�<��{?#Hd?��?���?fm?C��>F@e?[��=��1?M>`H�?T�?���>��>Z�b?%��>$�?c�<um*?�(?3*?��?D@?{H�>]��>k�E?9�v=;�:?q-<?�

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

@@ -1 +1 @@
-
B
yJ�
ɚ�?���>%�]?,��?(��?�~]��{X?ng��ӽ'��>��>���?A�3?b��=�c�>�ѧ><��?��P�
+
B
yJ�
ɚ�?���>%�]?,��?)��?�~]��{X?ng��ӽ'��>��>���?A�3?b��=�c�>�ѧ><��?��P�

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

@@ -1 +1 @@
-
B
yJ�
#
+
B
yJ�
#

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

@@ -1,2 +1,2 @@
-
B
yJ�
��?q�e?̍2?eo?�~�>+�D?�@�>B�?�W�?���>���?�?�%?/9�?0��=��=զ�<�*�?5�?���?C�@�R�?���>��?w]�=�B?��>
-��??1?���>���>��?
!�>;,%?��<Ξ8?W6?�T8?ܸ�?;U?��>���>��\?��v=�
+
B
yJ�
��?q�e?̍2?eo?�~�>+�D?�@�>B�?�W�?���>���?�?�%?/9�?0��=��=Ԧ�<�*�?5�?���?C�@�R�?���>��?w]�=�B?��>
+��??1?���>���>��?
!�>;,%?��<Ϟ8?W6?�T8?۸�?;U?��>���>��\?��v=�

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|>