onnx 1.17.0__cp310-cp310-win_amd64.whl → 1.18.0__cp310-cp310-win_amd64.whl

onnx/reference/ops/experimental/_op_list.py

@@ -6,14 +6,13 @@ from __future__ import annotations
 import textwrap
 from typing import Any
 
-from onnx.reference.op_run import OpFunction
+from onnx.reference.op_run import OpFunction, OpRun
 from onnx.reference.ops._helpers import build_registered_operators_any_domain
-from onnx.reference.ops.experimental._op_run_experimental import OpRunExperimental
 from onnx.reference.ops.experimental.op_im2col import Im2Col  # noqa: F401
 
 
-def _build_registered_operators() -> dict[str, dict[int | None, OpRunExperimental]]:
-    return build_registered_operators_any_domain(globals().copy())  # type: ignore[return-value]
+def _build_registered_operators() -> dict[str, dict[int | None, type[OpRun]]]:
+    return build_registered_operators_any_domain(globals().copy())
 
 
 def load_op(
@@ -21,7 +20,6 @@ def load_op(
     op_type: str,
     version: None | int,
     custom: Any = None,
-    evaluator_cls: type | None = None,  # noqa: ARG001
 ) -> Any:
     """Loads the implemented for a specified operator.
 
@@ -30,25 +28,24 @@ def load_op(
         op_type: oprator type
         version: requested version
         custom: custom implementation (like a function)
-        evaluator_cls: unused
 
     Returns:
         class
     """
     global _registered_operators  # noqa: PLW0603
     if _registered_operators is None:
-        _registered_operators = _build_registered_operators()  # type: ignore[assignment]
+        _registered_operators = _build_registered_operators()
     if custom is not None:
-        return lambda *args: OpFunction(*args, impl=custom)  # type: ignore
+        return lambda *args: OpFunction(*args, impl=custom)
     if domain != "experimental":
         raise ValueError(f"Domain must be '' not {domain!r}.")
-    if op_type not in _registered_operators:  # type: ignore
+    if op_type not in _registered_operators:
         available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))  # type: ignore
         raise NotImplementedError(
             f"No registered implementation for operator {op_type!r} "
             f"and domain {domain!r} in\n{available}"
         )
-    impl = _registered_operators[op_type]  # type: ignore
+    impl = _registered_operators[op_type]
     if None not in impl:
         raise RuntimeError(
             f"No default implementation for operator {op_type!r} "
@@ -56,28 +53,20 @@ def load_op(
             f"{', '.join(map(str, impl))}."
         )
     if version is None or len(impl) == 1:
-        cl = impl[None]
-    else:
-        best = -1
-        for v in impl:
-            if v is None:
-                continue
-            if best < v <= version:
-                best = v
-        if best == -1:
-            raise RuntimeError(
-                f"No implementation for operator {op_type!r} "
-                f"domain {domain!r} and version {version!r}, found "
-                f"{', '.join(map(str, impl))}."
-            )
-        cl = impl[best]
-    if cl is None:
-        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))  # type: ignore
-        raise ValueError(
-            f"Not registered implementation for operator {op_type!r}, "
-            f"domain {domain!r}, and {version!r} in\n{available}"
+        return impl[None]
+    best = -1
+    for v in impl:
+        if v is None:
+            continue
+        if best < v <= version:
+            best = v
+    if best == -1:
+        raise RuntimeError(
+            f"No implementation for operator {op_type!r} "
+            f"domain {domain!r} and version {version!r}, found "
+            f"{', '.join(map(str, impl))}."
         )
-    return cl
+    return impl[best]
 
 
-_registered_operators: dict[str, dict[int | None, OpRunExperimental]] | None = None
+_registered_operators: dict[str, dict[int | None, type[OpRun]]] | None = None

onnx/reference/ops/op_affine_grid.py

@@ -46,9 +46,9 @@ def construct_original_grid(data_size, align_corners):
 def apply_affine_transform(theta_n, original_grid_homo):
     # theta_n: (N, 2, 3) for 2D, (N, 3, 4) for 3D
     # original_grid_homo: (H, W, 3) for 2D, (D, H, W, 4) for 3D
-    assert (
-        theta_n.ndim == 3
-    ), "theta_n shall have shape of (N, 2, 3) for 2D, (N, 3, 4) for 3D"
+    assert theta_n.ndim == 3, (
+        "theta_n shall have shape of (N, 2, 3) for 2D, (N, 3, 4) for 3D"
+    )
     if original_grid_homo.ndim == 3:
         N, dim_2d, dim_homo = theta_n.shape
         assert dim_2d == 2 and dim_homo == 3

onnx/reference/ops/op_attention.py

@@ -0,0 +1,211 @@
+# Copyright (c) ONNX Project Contributors
+
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import numpy as np
+
+import onnx
+from onnx.reference.op_run import OpRun
+
+
+def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
+    x_max = np.max(x, axis=axis, keepdims=True)
+    tmp = np.exp(x - x_max)
+    s = np.sum(tmp, axis=axis, keepdims=True)
+    return tmp / s
+
+
+def _softcap(X, softcap):
+    if softcap > 0:
+        Y = X / softcap
+        Y = np.tanh(Y)
+        return Y * softcap
+    else:
+        return X
+
+
+def _compute_attention(
+    Q: np.ndarray,
+    K: np.ndarray,
+    V: np.ndarray,
+    attn_mask: np.ndarray | None = None,
+    past_key: np.ndarray | None = None,
+    past_value: np.ndarray | None = None,
+    scale=None,
+    is_causal=False,
+    q_num_heads=None,
+    kv_num_heads=None,
+    softmax_precision=None,
+    softcap=None,
+    qk_matmul_output_mode=None,
+) -> np.ndarray:
+    assert len(Q.shape) == len(K.shape) == len(V.shape)
+    # Set input tensors (Q, K, V) to the correct shape if input shape is 3D
+    # NewShapeQ (batch_size, q_num_heads, q_sequence_length, head_size)
+    # NewShapeK  (batch_size, kv_num_heads, kv_sequence_length, head_size)
+    # NewShapeV (value) has shape (batch_size, kv_num_heads, kv_sequence_length, v_head_size)
+    input_shape_len = len(Q.shape)
+    batch_size = Q.shape[0]
+    if len(Q.shape) == 3:
+        hidden_size_q = Q.shape[2]
+        hidden_size_k = K.shape[2]
+        hidden_size_v = V.shape[2]
+        assert q_num_heads is not None and kv_num_heads is not None
+
+        head_size_q = int(hidden_size_q / q_num_heads)
+        new_shape_q = [batch_size, q_num_heads, Q.shape[1], head_size_q]
+        Q = np.reshape(Q, new_shape_q)
+
+        head_size_k = int(hidden_size_k / kv_num_heads)
+        new_shape_k = [batch_size, kv_num_heads, K.shape[1], head_size_k]
+        K = np.reshape(K, new_shape_k)
+
+        head_size_v = int(hidden_size_v / kv_num_heads)
+        new_shape_v = [batch_size, kv_num_heads, V.shape[1], head_size_v]
+        V = np.reshape(V, new_shape_v)
+    assert len(Q.shape) == 4 and len(K.shape) == 4 and len(V.shape) == 4
+
+    # Calculate Scaling Factor if not provided
+    if scale is None:
+        q_head_size = Q.shape[3]
+        scale = 1 / np.sqrt(q_head_size)
+    scale = np.sqrt(scale)
+
+    # Update key and value cache
+    if past_key is not None:
+        present_key = np.concatenate((past_key, K), axis=2)
+    else:
+        present_key = K
+    if past_value is not None:
+        present_value = np.concatenate((past_value, V), axis=2)
+    else:
+        present_value = V
+    K = present_key
+    V = present_value
+
+    # Create attn_bias
+    q_sequence_length = Q.shape[2]
+    kv_sequence_length = K.shape[2]
+    attn_bias = np.zeros((q_sequence_length, kv_sequence_length), dtype=Q.dtype)
+    # First case: If is_causal is provided
+    # If set to true, the attention masking is a lower triangular matrix when the mask
+    # is a square matrix. The attention masking has the form of the upper left causal
+    # bias due to the alignment when the mask is a non-square matrix.
+    if is_causal == 1:
+        assert attn_mask is None
+        temp_mask = np.ones((q_sequence_length, kv_sequence_length), dtype=bool)
+        temp_mask = np.tril(temp_mask, k=0)
+        temp_mask = np.logical_not(temp_mask)
+        attn_bias_ma = np.ma.array(attn_bias, mask=temp_mask)
+        attn_bias = attn_bias_ma.filled(fill_value=float("-inf"))
+    if attn_mask is not None:
+        assert is_causal != 1
+        if attn_mask.dtype == bool:
+            attn_mask = np.logical_not(attn_mask)
+            attn_bias_ma = np.ma.array(attn_bias, mask=attn_mask)
+            attn_bias = attn_bias_ma.filled(fill_value=float("-inf"))
+        else:
+            attn_bias += attn_mask
+
+    # Group Query Attention is applied if the following are satisfied
+    # 1) q_num_heads != kv_num_heads
+    # 2) q_num_heads % kv_num_heads == 0
+    # 3) kv_num_heads == k_num_heads == v_num_heads
+    if q_num_heads is None:
+        q_num_heads = Q.shape[1]
+    if kv_num_heads is None:
+        k_num_heads = K.shape[1]
+        v_num_heads = K.shape[1]
+    else:
+        k_num_heads = kv_num_heads
+        v_num_heads = kv_num_heads
+    if (
+        (q_num_heads != k_num_heads)
+        and (q_num_heads % k_num_heads == 0)
+        and (k_num_heads == v_num_heads)
+    ):
+        seq_reps = int(q_num_heads / k_num_heads)
+        reps = [1, seq_reps, 1, 1]
+        K = np.tile(K, reps)
+        V = np.tile(V, reps)
+
+    # The following pattern is applied
+    #      Q          K          V
+    #      |          |          |
+    #     Q*scale    K*scale     |
+    #      |          |          |
+    #      |       Transpose     |
+    #      |          |          |
+    #      ---MatMul---          |
+    #            |               |
+    # at_mask---Add              |
+    #            |               |
+    #  softcap (if provided)     |
+    #            |               |
+    #         Softmax            |
+    #            |               |
+    #            -----MatMul------
+    #                    |
+    #                    Y
+    k_transpose = np.transpose(K, (0, 1, 3, 2))
+    qk_matmul_output = np.matmul(Q * scale, k_transpose * scale)
+    qk_with_bias = qk_matmul_output + attn_bias
+    if qk_matmul_output_mode == 1:
+        qk_matmul_output = qk_matmul_output + attn_bias
+
+    # Apply softcap
+    if softcap is not None:
+        qk_with_bias = _softcap(qk_with_bias, softcap)
+        if qk_matmul_output_mode == 2:
+            qk_matmul_output = qk_with_bias
+
+    if softmax_precision is not None:
+        qk_with_bias = qk_with_bias.astype(
+            onnx.helper.tensor_dtype_to_np_dtype(softmax_precision)
+        )
+    qk_softmax = _softmax(qk_with_bias)
+    if qk_matmul_output_mode == 3:
+        qk_matmul_output = qk_softmax
+    qk_matmul_output = qk_matmul_output.astype(Q.dtype)
+
+    output = np.matmul(qk_softmax, V).astype(Q.dtype)
+    if input_shape_len == 3:
+        output = np.transpose(output, (0, 2, 1, 3))
+        output = np.reshape(output, (output.shape[0], output.shape[1], -1))
+    return output, present_key, present_value, qk_matmul_output
+
+
+class Attention(OpRun):
+    def _run(
+        self,
+        Q: np.ndarray,
+        K: np.ndarray,
+        V: np.ndarray,
+        attn_mask: np.ndarray | None = None,
+        past_key: np.ndarray | None = None,
+        past_value: np.ndarray | None = None,
+        scale=None,
+        is_causal=False,
+        q_num_heads=None,
+        kv_num_heads=None,
+        softmax_precision=None,
+        softcap=None,
+        qk_matmul_output_mode=None,
+    ) -> np.ndarray:
+        res = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+            scale=scale,
+            is_causal=is_causal,
+            q_num_heads=q_num_heads,
+            kv_num_heads=kv_num_heads,
+            softmax_precision=softmax_precision,
+            softcap=softcap,
+            qk_matmul_output_mode=qk_matmul_output_mode,
+        )
+        return res

onnx/reference/ops/op_cast.py

@@ -8,6 +8,7 @@ import numpy as np
 from onnx import subbyte
 from onnx._custom_element_types import (
     bfloat16,
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -25,6 +26,7 @@ from onnx.numpy_helper import (
     bfloat16_to_float32,
     float8e4m3_to_float32,
     float8e5m2_to_float32,
+    unpacked_float4e2m1_to_float32,
 )
 from onnx.onnx_pb import TensorProto
 from onnx.reference.op_run import OpRun
@@ -126,6 +128,20 @@ def cast_to(x, to, saturate):  # noqa: PLR0911
                 y[i] = el
             return y.reshape(x.shape)
 
+    if x.dtype == float4e2m1 and x.dtype.descr[0][0] == "float4e2m1":
+        if to == TensorProto.FLOAT4E2M1:
+            return x
+        res = unpacked_float4e2m1_to_float32(x)
+        if to == TensorProto.FLOAT:
+            return res.astype(np.float32)
+        elif to == TensorProto.FLOAT16:
+            return res.astype(np.float16)
+
+    if to == TensorProto.FLOAT4E2M1:
+        xf = x.astype(np.float32)
+        y = subbyte.float32_to_float4e2m1_unpacked(xf).astype(float4e2m1)
+        return y.reshape(x.shape)
+
     if to == TensorProto.STRING:
         return x.astype(np.str_)
 

onnx/reference/ops/op_cast_like.py

@@ -9,6 +9,7 @@ from onnx.reference.op_run import OpRun
 from onnx.reference.ops.op_cast import (
     bfloat16,
     cast_to,
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -34,6 +35,8 @@ def _cast_like(x, y, saturate):
         to = TensorProto.UINT4
     elif y.dtype == int4 and y.dtype.descr[0][0] == "int4":
         to = TensorProto.INT4
+    elif y.dtype == float4e2m1 and y.dtype.descr[0][0] == "float4e2m1":
+        to = TensorProto.FLOAT4E2M1
     else:
         to = np_dtype_to_tensor_dtype(y.dtype)  # type: ignore
     return (cast_to(x, to, saturate),)

onnx/reference/ops/op_concat_from_sequence.py

@@ -12,8 +12,12 @@ from onnx.reference.op_run import OpRun
 
 def _concat_from_sequence(seq: list[Any], axis: int, new_axis: int = 0) -> np.ndarray:
     if new_axis == 1:
-        seq2 = [s[..., np.newaxis] for s in seq]
-        res = np.concatenate(seq2, axis=-1)
+        if axis == -1:
+            seq2 = [s[..., np.newaxis] for s in seq]
+            res = np.concatenate(seq2, axis=-1)
+        else:
+            seq2 = [np.expand_dims(s, axis) for s in seq]
+            res = np.concatenate(seq2, axis=axis)
     else:
         res = np.concatenate(seq, axis=axis)
     return res  # type: ignore

onnx/reference/ops/op_constant.py

@@ -7,6 +7,7 @@ import numpy as np
 
 from onnx._custom_element_types import (
     bfloat16,
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -27,6 +28,7 @@ def _check_dtype(val):  # type: ignore
         float8e5m2fnuz,
         uint4,
         int4,
+        float4e2m1,
         np.int8,
         np.uint8,
         np.float16,
@@ -108,7 +110,7 @@ class Constant_12(ConstantCommon):
             self.cst_convert = lambda v: v
         elif hasattr(self, "value") and self.value is not None:  # type: ignore
             self.cst_name = "value"  # type: ignore
-            self.cst = self.value if isinstance(self.value, RefAttrName) else self.value  # type: ignore
+            self.cst = self.value  # type: ignore
             self.cst_convert = lambda v: v
         else:
             for attr, np_dtype in {

onnx/reference/ops/op_conv.py

@@ -276,15 +276,11 @@ def _conv_implementation(  # type: ignore
                                         )
                                     s = np.dot(
                                         img.reshape((1, -1)), w_.reshape((-1, 1))
-                                    )[
-                                        0, 0
-                                    ]  # (img * w_).sum()
+                                    )[0, 0]  # (img * w_).sum()
                                 else:
                                     s = np.dot(
                                         img.reshape((1, -1)), w.reshape((-1, 1))
-                                    )[
-                                        0, 0
-                                    ]  # (img * w).sum()
+                                    )[0, 0]  # (img * w).sum()
                                 res[n, nw, hr, wr, zr] += s  # type: ignore
 
         return res

onnx/reference/ops/op_conv_transpose.py

@@ -107,9 +107,7 @@ class ConvTranspose(OpRun):
             for group_id in range(group):
                 group_X = X[:, group_id * C // group : (group_id + 1) * C // group, ...]
                 group_W = W[
-                    group_id
-                    * num_output_channels
-                    // group : (group_id + 1)
+                    group_id * num_output_channels // group : (group_id + 1)
                     * num_output_channels
                     // group,
                     ...,

onnx/reference/ops/op_cum_sum.py

@@ -9,12 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 class CumSum(OpRun):
-    def _run(self, x, *axis, exclusive=None, reverse=None):  # type: ignore
-        axis = None if not axis else axis[0]  # type: ignore
-        if axis is None:  # type: ignore
-            if reverse or exclusive:
-                raise NotImplementedError("reverse=1 or exclusive=1 not implemented")
-            return (np.cumsum(x),)
+    def _run(self, x, axis, exclusive=None, reverse=None):  # type: ignore
         if not isinstance(axis, (np.int32, np.int64)):
             if len(axis.shape) > 1 or (len(axis.shape) > 0 and axis.shape[0] != 1):  # type: ignore
                 raise RuntimeError(
@@ -34,7 +29,7 @@ class CumSum(OpRun):
             res = np.zeros(x.shape, dtype=x.dtype)
             np.cumsum(x[tuple(indices_c)], axis=axis, out=res[tuple(indices_d)])  # type: ignore
         else:
-            res = np.cumsum(x, axis=axis)  # type: ignore
+            res = np.cumsum(x, axis=axis, dtype=x.dtype)  # type: ignore
         if reverse:
             res = res[tuple(rev_indices)]
         return (res,)

onnx/reference/ops/op_dequantize_linear.py

@@ -7,6 +7,7 @@ import numpy as np
 
 from onnx import TensorProto
 from onnx._custom_element_types import (
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -14,8 +15,12 @@ from onnx._custom_element_types import (
     int4,
     uint4,
 )
-from onnx.helper import np_dtype_to_tensor_dtype
-from onnx.numpy_helper import float8e4m3_to_float32, float8e5m2_to_float32
+from onnx.helper import np_dtype_to_tensor_dtype, tensor_dtype_to_np_dtype
+from onnx.numpy_helper import (
+    float8e4m3_to_float32,
+    float8e5m2_to_float32,
+    unpacked_float4e2m1_to_float32,
+)
 from onnx.reference.op_run import OpRun
 from onnx.reference.ops.op_quantize_linear import reshape_input
 
@@ -35,6 +40,8 @@ class _CommonDequantizeLinear(OpRun):
             tensor_dtype = TensorProto.UINT4
         elif x.dtype == int4 and x.dtype.descr[0][0] == "int4":
             tensor_dtype = TensorProto.INT4
+        elif x.dtype == float4e2m1 and x.dtype.descr[0][0] == "float4e2m1":
+            tensor_dtype = TensorProto.FLOAT4E2M1
         else:
             tensor_dtype = np_dtype_to_tensor_dtype(x.dtype)
         return tensor_dtype
@@ -46,6 +53,7 @@ class _CommonDequantizeLinear(OpRun):
         x_zero_point: np.ndarray | None = None,
         axis: int | None = None,
         block_size: int | None = None,
+        output_dtype: int | None = None,
     ):  # type: ignore
         x_type = self.get_x_type(x)
         fp8_type = x_type in {
@@ -54,7 +62,11 @@ class _CommonDequantizeLinear(OpRun):
             TensorProto.FLOAT8E5M2,
             TensorProto.FLOAT8E5M2FNUZ,
         }
-        if x_zero_point is not None and not fp8_type:
+        if (
+            x_zero_point is not None
+            and not fp8_type
+            and x_type != TensorProto.FLOAT4E2M1
+        ):
             zero_type = self.get_x_type(x_zero_point)
             if x_type != zero_type:
                 raise ValueError(
@@ -81,10 +93,18 @@ class _CommonDequantizeLinear(OpRun):
                 dx = float8e5m2_to_float32(x)
             elif x_type == TensorProto.FLOAT8E5M2FNUZ:
                 dx = float8e5m2_to_float32(x, fn=True, uz=True)
+            elif x_type == TensorProto.FLOAT4E2M1:
+                dx = unpacked_float4e2m1_to_float32(x)
             else:
                 dx = x.astype(np.float32)
         y = dx * reshape_input(x_scale, x.shape, axis, block_size)
-        return (y.astype(x_scale.dtype),)
+        return (
+            y.astype(
+                tensor_dtype_to_np_dtype(output_dtype)
+                if output_dtype
+                else x_scale.dtype
+            ),
+        )
 
 
 class DequantizeLinear_19(_CommonDequantizeLinear):
@@ -98,3 +118,11 @@ class DequantizeLinear_21(_CommonDequantizeLinear):
     def _run(self, *args, axis=None, block_size=None):  # type: ignore
         # args: x, y_scale, zero_point
         return super()._run(*args, axis=axis, block_size=block_size)  # type: ignore
+
+
+class DequantizeLinear_23(_CommonDequantizeLinear):
+    def _run(self, *args, axis=None, block_size=None, output_dtype=None):  # type: ignore
+        # args: x, y_scale, zero_point
+        return super()._run(
+            *args, axis=axis, block_size=block_size, output_dtype=output_dtype
+        )  # type: ignore

onnx/reference/ops/op_if.py

@@ -4,10 +4,13 @@
 
 from __future__ import annotations
 
-import numpy as np
+from typing import TYPE_CHECKING
 
 from onnx.reference.op_run import OpRun
 
+if TYPE_CHECKING:
+    import numpy as np
+
 
 class If(OpRun):
     def __init__(self, onnx_node, run_params):  # type: ignore

onnx/reference/ops/op_leaky_relu.py

@@ -3,10 +3,13 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-import numpy as np
+from typing import TYPE_CHECKING
 
 from onnx.reference.ops._op import OpRunUnaryNum
 
+if TYPE_CHECKING:
+    import numpy as np
+
 
 def _leaky_relu(x: np.ndarray, alpha: float) -> np.ndarray:
     sign = (x > 0).astype(x.dtype)

onnx/reference/ops/op_loop.py

@@ -45,7 +45,7 @@ class Loop(OpRun):
             begin = len(loop_inputs) - len(args)
             all_inputs = loop_inputs[begin:]
             for name, val in zip(all_inputs, args):
-                inputs[name] = val
+                inputs[name] = val  # noqa: PERF403
         if context is not None:
             for a in context:
                 inputs[a] = context[a]

onnx/reference/ops/op_non_max_suppression.py

@@ -25,7 +25,7 @@ class PrepareContext:
 
 
 class SelectedIndex:
-    __slots__ = ("batch_index_", "class_index_", "box_index_")
+    __slots__ = ("batch_index_", "box_index_", "class_index_")
 
     def __init__(
         self, batch_index: int = 0, class_index: int = 0, box_index: int = 0
@@ -160,7 +160,7 @@ class NonMaxSuppression(OpRun):
             pc.max_output_boxes_per_class_ = max_output_boxes_per_class_tensor
         if iou_threshold_tensor.size != 0:
             pc.iou_threshold_ = iou_threshold_tensor
-        if score_threshold_tensor.size != 0:
+        if score_threshold_tensor is not None and score_threshold_tensor.size != 0:
             pc.score_threshold_ = score_threshold_tensor
 
         pc.boxes_size_ = boxes_tensor.size
@@ -177,13 +177,11 @@ class NonMaxSuppression(OpRun):
         self,
         boxes,
         scores,
-        max_output_boxes_per_class,
-        iou_threshold,
-        score_threshold,
-        center_point_box,
+        max_output_boxes_per_class=None,
+        iou_threshold=None,
+        score_threshold=None,
+        center_point_box=None,
     ):
-        center_point_box = center_point_box or self.center_point_box  # type: ignore
-
         pc = PrepareContext()
         self.prepare_compute(
             pc,