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

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

@@ -3,22 +3,23 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-import sys
+import itertools
 
 import numpy as np
 
 import onnx
-from onnx import TensorProto, helper
+from onnx import TensorProto
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
-from onnx.helper import float32_to_float8e4m3, float32_to_float8e5m2, make_tensor
-from onnx.numpy_helper import float8e4m3_to_float32, float8e5m2_to_float32
+from onnx.helper import make_tensor, tensor_dtype_to_np_dtype
+
+F8_TYPES = frozenset({"FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ"})
+FOUR_BIT_TYPES = frozenset({"UINT4", "INT4", "FLOAT4E2M1"})
 
 
 class CastLike(Base):
     @staticmethod
     def export() -> None:
-        shape = (3, 4)
         test_cases = [
             ("FLOAT", "FLOAT16"),
             ("FLOAT", "DOUBLE"),
@@ -26,26 +27,51 @@ class CastLike(Base):
             ("FLOAT16", "DOUBLE"),
             ("DOUBLE", "FLOAT"),
             ("DOUBLE", "FLOAT16"),
-            ("FLOAT", "STRING"),
-            ("STRING", "FLOAT"),
             ("FLOAT", "BFLOAT16"),
             ("BFLOAT16", "FLOAT"),
             ("FLOAT", "FLOAT8E4M3FN"),
+            ("FLOAT16", "FLOAT8E4M3FN"),
             ("FLOAT", "FLOAT8E4M3FNUZ"),
+            ("FLOAT16", "FLOAT8E4M3FNUZ"),
             ("FLOAT8E4M3FN", "FLOAT"),
+            ("FLOAT8E4M3FN", "FLOAT16"),
             ("FLOAT8E4M3FNUZ", "FLOAT"),
+            ("FLOAT8E4M3FNUZ", "FLOAT16"),
             ("FLOAT", "FLOAT8E5M2"),
+            ("FLOAT16", "FLOAT8E5M2"),
             ("FLOAT", "FLOAT8E5M2FNUZ"),
+            ("FLOAT16", "FLOAT8E5M2FNUZ"),
             ("FLOAT8E5M2", "FLOAT"),
+            ("FLOAT8E5M2", "FLOAT16"),
             ("FLOAT8E5M2FNUZ", "FLOAT"),
+            ("FLOAT8E5M2FNUZ", "FLOAT16"),
+            ("FLOAT", "UINT4"),
+            ("FLOAT16", "UINT4"),
+            ("FLOAT", "INT4"),
+            ("FLOAT16", "INT4"),
+            ("UINT4", "FLOAT"),
+            ("UINT4", "FLOAT16"),
+            ("UINT4", "UINT8"),
+            ("INT4", "FLOAT"),
+            ("INT4", "FLOAT16"),
+            ("INT4", "INT8"),
+            ("FLOAT4E2M1", "FLOAT"),
+            ("FLOAT4E2M1", "FLOAT16"),
+            ("FLOAT", "FLOAT4E2M1"),
+            ("FLOAT16", "FLOAT4E2M1"),
         ]
 
-        vect_float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
-        vect_float32_to_float8e5m2 = np.vectorize(float32_to_float8e5m2)
+        f8_types = {"FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ"}
 
         for from_type, to_type in test_cases:
-            input_type_proto = None
-            output_type_proto = None
+            if from_type == to_type:
+                # Skip cases where from_type and to_type are the same
+                continue
+            from_dtype = getattr(TensorProto, from_type)
+            to_dtype = getattr(TensorProto, to_type)
+            from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
+            to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)
+
             if from_type == "BFLOAT16" or to_type == "BFLOAT16":
                 np_fp32 = np.array(
                     [
@@ -64,50 +90,9 @@ class CastLike(Base):
                     ],
                     dtype=np.float32,
                 )
-                little_endisan = sys.byteorder == "little"
-                np_uint16_view = np_fp32.view(dtype=np.uint16)
-                np_bfp16 = (
-                    np_uint16_view[1::2] if little_endisan else np_uint16_view[0::2]
-                )
-                if to_type == "BFLOAT16":
-                    assert from_type == "FLOAT"
-                    input = np_fp32.reshape([3, 4])
-                    output = np_bfp16.reshape([3, 4])
-                    input_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.FLOAT), input.shape
-                    )
-                    output_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.BFLOAT16), output.shape
-                    )
-                else:
-                    assert to_type == "FLOAT"
-                    input = np_bfp16.reshape([3, 4])
-                    # convert bfloat to FLOAT
-                    np_fp32_zeros = np.zeros((len(np_bfp16) * 2,), dtype=np.uint16)
-                    if little_endisan:
-                        np_fp32_zeros[1::2] = np_bfp16
-                    else:
-                        np_fp32_zeros[0::2] = np_bfp16
-                    np_fp32_from_bfloat = np_fp32_zeros.view(dtype=np.float32)
-                    output = np_fp32_from_bfloat.reshape([3, 4])
-                    input_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.BFLOAT16), input.shape
-                    )
-                    output_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.FLOAT), output.shape
-                    )
-                like = output.flatten()[0:1]
-            elif from_type in (
-                "FLOAT8E4M3FN",
-                "FLOAT8E4M3FNUZ",
-                "FLOAT8E5M2",
-                "FLOAT8E5M2FNUZ",
-            ) or to_type in (
-                "FLOAT8E4M3FN",
-                "FLOAT8E4M3FNUZ",
-                "FLOAT8E5M2",
-                "FLOAT8E5M2FNUZ",
-            ):
+                input_shape = (3, 4)
+
+            elif from_type in f8_types or to_type in f8_types:
                 np_fp32 = np.array(
                     [
                         "0.47892547",
@@ -115,87 +100,48 @@ class CastLike(Base):
                         "0.49968487",
                         "0.81910545",
                         "0.47031248",
-                        "0.816468",
-                        "0.21087195",
                         "0.7229038",
+                        "1000000",
+                        "1e-7",
                         "NaN",
                         "INF",
                         "+INF",
                         "-INF",
+                        "-0.0000001",
+                        "0.0000001",
+                        "-1000000",
                     ],
                     dtype=np.float32,
                 )
-                if to_type == "FLOAT8E4M3FN":
-                    expected = float8e4m3_to_float32(
-                        vect_float32_to_float8e4m3(np_fp32)
-                    )
-                    expected_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FN, [3, 4], expected.tolist()
-                    )
-                    like_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FN, [1], expected[:1]
-                    )
-                elif to_type == "FLOAT8E4M3FNUZ":
-                    expected = float8e4m3_to_float32(
-                        vect_float32_to_float8e4m3(np_fp32, uz=True), uz=True
-                    )
-                    expected_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FNUZ, [3, 4], expected.tolist()
-                    )
-                    like_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FNUZ, [1], expected[:1]
-                    )
-                elif to_type == "FLOAT8E5M2":
-                    expected = float8e5m2_to_float32(
-                        vect_float32_to_float8e5m2(np_fp32)
-                    )
-                    expected_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2, [3, 4], expected.tolist()
-                    )
-                    like_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2, [1], expected[:1]
-                    )
-                elif to_type == "FLOAT8E5M2FNUZ":
-                    expected = float8e5m2_to_float32(
-                        vect_float32_to_float8e5m2(np_fp32, fn=True, uz=True),
-                        fn=True,
-                        uz=True,
-                    )
-                    expected_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2FNUZ, [3, 4], expected.tolist()
-                    )
-                    like_tensor = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2FNUZ, [1], expected[:1]
-                    )
-                if from_type == "FLOAT":
-                    input = np_fp32.reshape((3, 4))
-                    output = expected_tensor
-                    like = like_tensor
-                else:
-                    assert to_type == "FLOAT"
-                    input = expected_tensor
-                    output = expected.reshape((3, 4))
-                    like = output.flatten()[:1]
-            elif from_type != "STRING":
-                input = np.random.random_sample(shape).astype(
-                    helper.tensor_dtype_to_np_dtype(getattr(TensorProto, from_type))
+                input_shape = (3, 5)
+            elif from_type in ("UINT4", "INT4") or to_type in ("UINT4", "INT4"):
+                np_fp32 = np.arange(-9, 16).astype(np.float32)
+                input_shape = (5, 5)
+            elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
+                np_fp32 = np.array(
+                    [
+                        "0.48",
+                        "0.25",
+                        "1.05",
+                        "-3.5",
+                        "-8",
+                        "9",
+                        "1000000",
+                        "1e-7",
+                        "NaN",
+                        "INF",
+                        "+INF",
+                        "-INF",
+                        "-4",
+                        "0.01",
+                        "-0.0",
+                    ],
+                    dtype=np.float32,
                 )
-                if to_type == "STRING":
-                    # Converting input to str, then give it object dtype for generating script
-                    ss = []
-                    for i in input.flatten():
-                        s = str(i).encode("utf-8")
-                        su = s.decode("utf-8")
-                        ss.append(su)
-
-                    output = np.array(ss).astype(object).reshape([3, 4])
-                else:
-                    output = input.astype(
-                        helper.tensor_dtype_to_np_dtype(getattr(TensorProto, to_type))
-                    )
-                like = output.flatten()[0:1]
+                input_shape = (3, 5)
+
             else:
-                input = np.array(
+                np_fp32 = np.array(
                     [
                         "0.47892547",
                         "0.48033667",
@@ -210,34 +156,136 @@ class CastLike(Base):
                         "+INF",
                         "-INF",
                     ],
-                    dtype=np.dtype(object),
+                    dtype=np.float32,
                 ).reshape([3, 4])
-                output = input.astype(
-                    helper.tensor_dtype_to_np_dtype(getattr(TensorProto, to_type))
+                input_shape = (3, 4)
+
+            if from_type in F8_TYPES:
+                np_from = onnx.numpy_helper.saturate_cast(np_fp32, from_np_dtype)
+                input = make_tensor(
+                    "input",
+                    from_dtype,
+                    input_shape,
+                    vals=np_from,
+                    raw=True,
+                )
+            elif from_type in FOUR_BIT_TYPES:
+                np_from = np_fp32.astype(from_np_dtype)
+                packed = onnx.numpy_helper._pack_4bitx2(np_from)
+                input = make_tensor(
+                    "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
+                )
+            else:
+                np_from = np_fp32.astype(from_np_dtype)
+                input = make_tensor(
+                    "input", from_dtype, input_shape, vals=np_from, raw=True
+                )
+
+            if to_type in F8_TYPES:
+                output = make_tensor(
+                    "output",
+                    to_dtype,
+                    input_shape,
+                    vals=onnx.numpy_helper.saturate_cast(np_from, to_np_dtype),
+                    raw=True,
                 )
-                like = output.flatten()[0:1]
+            elif to_type in FOUR_BIT_TYPES:
+                packed = onnx.numpy_helper._pack_4bitx2(np_from.astype(to_np_dtype))
+                output = make_tensor(
+                    "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
+                )
+            else:
+                output = make_tensor(
+                    "output",
+                    to_dtype,
+                    input_shape,
+                    vals=np_from.astype(to_np_dtype),
+                    raw=True,
+                )
+
+            like = make_tensor("like", to_dtype, (0,), vals=[])
+
             node = onnx.helper.make_node(
                 "CastLike",
                 inputs=["input", "like"],
                 outputs=["output"],
             )
-            if input_type_proto and output_type_proto:
-                like_type_proto = onnx.helper.make_tensor_type_proto(
-                    output_type_proto.tensor_type.elem_type, like.shape
-                )
 
-                expect(
-                    node,
-                    inputs=[input, like],
-                    outputs=[output],
-                    name="test_castlike_" + from_type + "_to_" + to_type,
-                    input_type_protos=[input_type_proto, like_type_proto],
-                    output_type_protos=[output_type_proto],
-                )
-            else:
-                expect(
-                    node,
-                    inputs=[input, like],
-                    outputs=[output],
-                    name="test_castlike_" + from_type + "_to_" + to_type,
-                )
+            expect(
+                node,
+                inputs=[input, like],
+                outputs=[output],
+                name="test_castlike_" + from_type + "_to_" + to_type,
+            )
+
+    @staticmethod
+    def export_saturate_false() -> None:
+        test_cases = itertools.product(
+            [
+                "FLOAT",
+                "FLOAT16",
+            ],
+            [
+                "FLOAT8E4M3FN",
+                "FLOAT8E4M3FNUZ",
+                "FLOAT8E5M2",
+                "FLOAT8E5M2FNUZ",
+            ],
+        )
+        input_shape = (3, 5)
+        for from_type, to_type in test_cases:
+            from_dtype = getattr(TensorProto, from_type)
+            to_dtype = getattr(TensorProto, to_type)
+            from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
+            to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)
+            np_fp32 = np.array(
+                [
+                    "0.47892547",
+                    "0.48033667",
+                    "0.49968487",
+                    "0.81910545",
+                    "0.47031248",
+                    "0.7229038",
+                    "1000000",
+                    "1e-7",
+                    "NaN",
+                    "INF",
+                    "+INF",
+                    "-INF",
+                    "-0.0000001",
+                    "0.0000001",
+                    "-1000000",
+                ],
+                dtype=np.float32,
+            )
+
+            input = make_tensor(
+                "input",
+                from_dtype,
+                input_shape,
+                vals=np_fp32.astype(from_np_dtype),
+                raw=True,
+            )
+            output = make_tensor(
+                "output",
+                to_dtype,
+                input_shape,
+                vals=np_fp32.astype(from_np_dtype).astype(to_np_dtype),
+                raw=True,
+            )
+
+            like = make_tensor("like", to_dtype, (0,), vals=[])
+
+            node = onnx.helper.make_node(
+                "CastLike",
+                inputs=["input", "like"],
+                outputs=["output"],
+                saturate=0,
+            )
+
+            expect(
+                node,
+                inputs=[input, like],
+                outputs=[output],
+                name="test_castlike_no_saturate_" + from_type + "_to_" + to_type,
+            )

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

@@ -60,6 +60,17 @@ class Clip(Base):
             name="test_clip_splitbounds",
         )
 
+        x = np.array([-2, 0, 6]).astype(np.float32)
+        y = np.array([1, 1, 1]).astype(np.float32)
+        min_val = np.float32(2)
+        max_val = np.float32(1)
+        expect(
+            node,
+            inputs=[x, min_val, max_val],
+            outputs=[y],
+            name="test_clip_min_greater_than_max",
+        )
+
     @staticmethod
     def export_clip_default() -> None:
         node = onnx.helper.make_node(

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

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -11,6 +11,9 @@ import onnx
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 class Concat(Base):
     @staticmethod

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

@@ -86,3 +86,27 @@ class CumSum(Base):
         axis = np.int32(-1)
         y = np.array([1.0, 3.0, 6.0, 4.0, 9.0, 15.0]).astype(np.float64).reshape((2, 3))
         expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_2d_negative_axis")
+
+    @staticmethod
+    def export_cumsum_2d_int32() -> None:
+        node = onnx.helper.make_node(
+            "CumSum",
+            inputs=["x", "axis"],
+            outputs=["y"],
+        )
+        x = np.array([1, 2, 3, 4, 5, 6]).astype(np.int32).reshape((2, 3))
+        axis = np.int32(0)
+        y = np.array([1, 2, 3, 5, 7, 9]).astype(np.int32).reshape((2, 3))
+        expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_2d_int32")
+
+    @staticmethod
+    def export_cumsum_1d_int32_exclusive() -> None:
+        node = onnx.helper.make_node(
+            "CumSum", inputs=["x", "axis"], outputs=["y"], exclusive=1
+        )
+        x = np.array([1, 2, 3, 4, 5]).astype(np.int32)
+        axis = np.int32(0)
+        y = np.array([0, 1, 3, 6, 10]).astype(np.int32)
+        expect(
+            node, inputs=[x, axis], outputs=[y], name="test_cumsum_1d_int32_exclusive"
+        )

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

@@ -235,6 +235,28 @@ class DequantizeLinear(Base):
             name="test_dequantizelinear_int4",
         )
 
+    @staticmethod
+    def export_float4e2m1() -> None:
+        node = onnx.helper.make_node(
+            "DequantizeLinear",
+            inputs=["x", "x_scale", "x_zero_point"],
+            outputs=["y"],
+            axis=0,
+        )
+
+        # scalar zero point and scale
+        x = make_tensor("x", TensorProto.FLOAT4E2M1, [5], [0, 1, -1, 1.5, -4])
+        x_scale = np.float32(2)
+        x_zero_point = make_tensor("x_zero_point", TensorProto.FLOAT4E2M1, (1,), [0])
+        y = np.array([0, 2, -2, 3, -8], dtype=np.float32)
+
+        expect(
+            node,
+            inputs=[x, x_scale, x_zero_point],
+            outputs=[y],
+            name="test_dequantizelinear_float4e2m1",
+        )
+
     @staticmethod
     def export_blocked() -> None:
         node = onnx.helper.make_node(

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

@@ -29,11 +29,36 @@ class Div(Base):
         z = x / y
         expect(node, inputs=[x, y], outputs=[z], name="test_div")
 
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.int8) + 1
+        z = x // y
+        expect(node, inputs=[x, y], outputs=[z], name="test_div_int8")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.int16) + 1
+        z = x // y
+        expect(node, inputs=[x, y], outputs=[z], name="test_div_int16")
+
         x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
         y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8) + 1
         z = x // y
         expect(node, inputs=[x, y], outputs=[z], name="test_div_uint8")
 
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16) + 1
+        z = x // y
+        expect(node, inputs=[x, y], outputs=[z], name="test_div_uint16")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32) + 1
+        z = x // y
+        expect(node, inputs=[x, y], outputs=[z], name="test_div_uint32")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64) + 1
+        z = x // y
+        expect(node, inputs=[x, y], outputs=[z], name="test_div_uint64")
+
     @staticmethod
     def export_div_broadcast() -> None:
         node = onnx.helper.make_node(

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

@@ -11,7 +11,7 @@ from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 
 
-def dropout(X, drop_probability=0.5, seed=0, training_mode=False, return_mask=False):  # type: ignore
+def dropout(X, drop_probability=0.5, seed=0, training_mode=False, return_mask=False):
     if drop_probability == 0 or training_mode is False:
         if return_mask is True:
             return X, np.ones(X.shape, dtype=bool)

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

@@ -79,3 +79,15 @@ class Einsum(Base):
         Z = einsum_reference_implementation(Eqn, (X, Y))
 
         expect(node, inputs=[X, Y], outputs=[Z], name="test_einsum_batch_matmul")
+
+    @staticmethod
+    def export_einsum_scalar() -> None:
+        Eqn = "->"
+        node = onnx.helper.make_node(
+            "Einsum", inputs=["x"], outputs=["y"], equation=Eqn
+        )
+
+        X = np.array(5.0)  # scalar input
+        Z = einsum_reference_implementation(Eqn, (X,))
+
+        expect(node, inputs=[X], outputs=[Z], name="test_einsum_scalar")

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

@@ -24,6 +24,36 @@ class Equal(Base):
         z = np.equal(x, y)
         expect(node, inputs=[x, y], outputs=[z], name="test_equal")
 
+        x = (np.random.randn(3, 4, 5) * 10).astype(np.int8)
+        y = (np.random.randn(3, 4, 5) * 10).astype(np.int8)
+        z = np.equal(x, y)
+        expect(node, inputs=[x, y], outputs=[z], name="test_equal_int8")
+
+        x = (np.random.randn(3, 4, 5) * 10).astype(np.int16)
+        y = (np.random.randn(3, 4, 5) * 10).astype(np.int16)
+        z = np.equal(x, y)
+        expect(node, inputs=[x, y], outputs=[z], name="test_equal_int16")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
+        z = np.equal(x, y)
+        expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint8")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
+        z = np.equal(x, y)
+        expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint16")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
+        z = np.equal(x, y)
+        expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint32")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
+        z = np.equal(x, y)
+        expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint64")
+
     @staticmethod
     def export_equal_broadcast() -> None:
         node = onnx.helper.make_node(

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

@@ -11,7 +11,7 @@ from onnx.backend.test.case.node import expect
 
 
 # The below GatherElements' numpy implementation is from https://stackoverflow.com/a/46204790/11767360
-def gather_elements(data, indices, axis=0):  # type: ignore
+def gather_elements(data, indices, axis=0):
     data_swaped = np.swapaxes(data, 0, axis)
     index_swaped = np.swapaxes(indices, 0, axis)
     gathered = np.choose(index_swaped, data_swaped, mode="wrap")

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

@@ -25,7 +25,7 @@ def gather_nd_impl(
     # The number of elements in the batch_dims for data/indice array
     batch_dims_size = 1
 
-    # Check the shape of indice and data are identicial for batch dims.
+    # Check the shape of indice and data are identical for batch dims.
     for i in range(batch_dims):
         batch_dims_shape.append(indices.shape[i])
         batch_dims_size *= indices.shape[i]
@@ -48,7 +48,7 @@ def gather_nd_impl(
     reshaped_indices = indices.reshape(batch_dims_size, -1, indices.shape[-1])
 
     # Flatten 'data' to array of shape (batch_dim_size, data.shape[batch_dimes:])
-    reshaped_data = data.reshape((batch_dims_size,) + data.shape[batch_dims:])
+    reshaped_data = data.reshape((batch_dims_size, *data.shape[batch_dims:]))
 
     # gather each scalar value from 'data'
     for batch_dim in range(reshaped_indices.shape[0]):