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

onnx/reference/ops/op_pool_common.py

@@ -5,12 +5,15 @@ from __future__ import annotations
 
 import itertools
 import math
-from typing import Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from onnx.reference.op_run import OpRun
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 def get_pad_shape(
     auto_pad: str,
@@ -51,7 +54,7 @@ def get_pad_with_auto_pad(auto_pad: str, pad_shape: Sequence[int]) -> Sequence[i
 
 
 def get_output_shape_explicit_padding(
-    pads: Sequence[int],
+    pads: Sequence[int] | None,
     input_spatial_shape: Sequence[int],
     kernel_spatial_shape: Sequence[int],
     strides_spatial: Sequence[int],
@@ -93,6 +96,11 @@ def get_output_shape_explicit_padding(
 
         if ceil_mode:
             output_spatial_shape[dim] = int(np.ceil(dim_size))
+            # NOTE: ensure that the last pooling starts inside the image
+            if (output_spatial_shape[dim] - 1) * strides_spatial[
+                dim
+            ] >= input_spatial_shape[dim] + pads[dim]:
+                output_spatial_shape[dim] -= 1
         else:
             output_spatial_shape[dim] = int(np.floor(dim_size))
 
@@ -152,13 +160,6 @@ def get_output_shape_auto_pad(
     return out_shape
 
 
-def lp_pool(x: np.array, p: int) -> float:
-    y = 0
-    for v in np.nditer(x):
-        y += abs(v) ** p
-    return y ** (1.0 / p)
-
-
 def pool(
     padded: np.ndarray,
     x_shape: Sequence[int],
@@ -166,6 +167,7 @@ def pool(
     strides: Sequence[int],
     out_shape: Sequence[int],
     pooling_type: str,
+    pads_required: Sequence[int] | None = None,
     pads: Sequence[int] | None = None,
     dilations: Sequence[int] | None = None,
     count_include_pad: int = 0,
@@ -178,6 +180,7 @@ def pool(
     strides: the strides
     out_shape: the shape of the output tensor
     pooling_type: the pooling type, can be "AVG", "LPPOOL", or "MAX"
+    pads_required: the required padding to make sure the sliding window does not go out-of-bound
     pads: the padding in an order of head_pad_1, head_pad_2, ..., tail_pad_1, tail_pad_2, ...
     dilations: the dilation
     count_include_pad: whether to include the padding in the calculation of average and lp pooling
@@ -187,25 +190,27 @@ def pool(
     y = np.zeros([x_shape[0], x_shape[1], *list(out_shape)], dtype=padded.dtype)
     if dilations is None:
         dilations = np.ones([spatial_size], dtype=np.int64)
+    if pads_required is None:
+        pads_required = np.zeros([spatial_size * 2], dtype=np.int64)
+    elif len(pads_required) == 1:
+        pads_required = pads_required * spatial_size * 2
     if pads is None:
         pads = np.zeros([spatial_size * 2], dtype=np.int64)
     elif len(pads) == 1:
         pads = pads * spatial_size * 2
     strides = strides or [1] * spatial_size
 
-    def lp_pool_p(x):
-        return lp_pool(x, p)
-
+    # Iterate all the possible sliding windows
     for shape in itertools.product(
-        range(x_shape[0]),
-        range(x_shape[1]),
+        range(x_shape[0]),  # e.g. dim=0: [0]
+        range(x_shape[1]),  # e.g. dim=1: [0, 1]
         *[
             range(
                 int(
                     (
                         x_shape[i + 2]
-                        + pads[i]
-                        + pads[i + spatial_size]
+                        + pads_required[i]
+                        + pads_required[i + spatial_size]
                         - (1 + (kernel[i] - 1) * dilations[i])
                     )
                     / strides[i]
@@ -222,24 +227,34 @@ def pool(
                 for i in list(
                     itertools.product(
                         *[
-                            range(
-                                strides[i] * shape[i + 2],
-                                strides[i] * shape[i + 2]
-                                + (1 + (kernel[i] - 1) * dilations[i]),
-                                dilations[i],
-                            )
+                            [
+                                pixel
+                                for pixel in range(
+                                    strides[i] * shape[i + 2],
+                                    strides[i] * shape[i + 2]
+                                    + (1 + (kernel[i] - 1) * dilations[i]),
+                                    dilations[i],
+                                )
+                                if pixel
+                                < x_shape[i + 2] + pads[i] + pads[spatial_size + i]
+                            ]
                             for i in range(spatial_size)
                         ]
                     )
                 )
             ]
         )
+
         if pooling_type == "AVG":
             f = np.average
         elif pooling_type == "MAX":
             f = np.max
         elif pooling_type == "LPPOOL":
-            f = lp_pool_p
+
+            def lp_pool(x: np.array, p: int = p) -> float:
+                return np.sum(np.abs(x) ** p) ** (1.0 / p)
+
+            f = lp_pool
         else:
             raise NotImplementedError(
                 f"Pooling type {pooling_type} does not support. Should be AVG, MAX"
@@ -270,9 +285,9 @@ class CommonPool(OpRun):
         pading_value = np.nan if pooling_type == "MAX" or count_include_pad == 0 else 0
 
         if auto_pad in ["SAME_UPPER", "SAME_LOWER", "VALID"]:
-            assert (
-                ceil_mode is None or ceil_mode == 0
-            ), "ceil_mode is not supported with auto_pad"
+            assert ceil_mode is None or ceil_mode == 0, (
+                "ceil_mode is not supported with auto_pad"
+            )
             out_shape = get_output_shape_auto_pad(
                 auto_pad, x.shape[2:], kernel_shape, strides
             )
@@ -296,18 +311,19 @@ class CommonPool(OpRun):
                 out_shape,
                 pooling_type,
                 pads,
+                pads,
                 dilations,
                 count_include_pad,
                 p,
             )
             return (y,)
         else:
-            out_shape, pads = get_output_shape_explicit_padding(
+            out_shape, extra_pads = get_output_shape_explicit_padding(
                 pads, x_shape[2:], kernel_shape, strides, dilations, ceil_mode
             )
             # convert pads from [x1_begin, x2_begin,...,x1_end, x2_end,...] to [(x1_begin, x1_end), (x2_begin, x2_end),...]
-            n_dims = len(pads) // 2
-            pads_np = [(pads[i], pads[i + n_dims]) for i in range(n_dims)]
+            n_dims = len(extra_pads) // 2
+            pads_np = [(extra_pads[i], extra_pads[i + n_dims]) for i in range(n_dims)]
             padded = np.pad(
                 x,
                 ((0, 0), (0, 0), *pads_np),
@@ -321,6 +337,7 @@ class CommonPool(OpRun):
                 strides,
                 out_shape,
                 pooling_type,
+                extra_pads,
                 pads,
                 dilations,
                 count_include_pad,

onnx/reference/ops/op_qlinear_conv.py

@@ -50,6 +50,19 @@ class QLinearConv(OpRun):
             X, W, B, auto_pad, dilations, group, kernel_shape, pads, strides
         ).astype(np.int32)
 
+        # w_scale could be a scalar or a 1-D tensor. A 1-D tensor means a per
+        # output channel quantization.
+        if np.size(w_scale) > 1:
+            if np.ndim(w_scale) != 1:
+                raise ValueError(
+                    f"w_scale must be a scalar or a 1-D tensor. Got shape {np.shape(w_scale)}."
+                )
+            if np.size(w_scale) != np.shape(w)[0]:
+                raise ValueError(
+                    f"w_scale elements must match output channels: {np.size(w_scale)} != {np.shape(w)[0]}"
+                )
+            w_scale = np.expand_dims(w_scale, (0, 2, 3))
+
         R = res * (x_scale * w_scale / y_scale)
         if y_zero_point is not None:
             R += y_zero_point

onnx/reference/ops/op_quantize_linear.py

@@ -11,6 +11,7 @@ import numpy as np
 
 from onnx import TensorProto, subbyte
 from onnx._custom_element_types import (
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -30,7 +31,7 @@ from onnx.reference.op_run import OpRun
 def reshape_input(
     value: np.ndarray,
     shape: tuple[int, ...],
-    axis: int | None = None,
+    axis: int,
     block_size: int | None = None,
 ) -> np.ndarray:
     """Reshape/Replicate scale/zero-point to be broadcastable to shape.
@@ -97,6 +98,7 @@ class _CommonQuantizeLinear(OpRun):
         TensorProto.FLOAT8E4M3FNUZ,
         TensorProto.FLOAT8E5M2,
         TensorProto.FLOAT8E5M2FNUZ,
+        TensorProto.FLOAT4E2M1,
     )
 
     def get_zero_point_type(self, zero_point: np.ndarray) -> int:
@@ -122,11 +124,16 @@ class _CommonQuantizeLinear(OpRun):
             zero_point_type = TensorProto.UINT4
         elif zero_point.dtype == int4 and zero_point.dtype.descr[0][0] == "int4":
             zero_point_type = TensorProto.INT4
+        elif (
+            zero_point.dtype == float4e2m1
+            and zero_point.dtype.descr[0][0] == "float4e2m1"
+        ):
+            zero_point_type = TensorProto.FLOAT4E2M1
         else:
             zero_point_type = np_dtype_to_tensor_dtype(zero_point.dtype)
         return zero_point_type
 
-    def _run(
+    def _run(  # noqa: PLR0911
         self,
         x: np.ndarray,
         y_scale: np.ndarray,
@@ -134,7 +141,8 @@ class _CommonQuantizeLinear(OpRun):
         axis: int = 1,
         saturate: bool = True,
         block_size: int | None = None,
-        output_dtype: np.dtype | None = None,
+        output_dtype: int | None = None,
+        precision: int | None = None,
     ) -> tuple[np.ndarray]:
         y_scale = reshape_input(y_scale, x.shape, axis, block_size)
 
@@ -160,7 +168,11 @@ class _CommonQuantizeLinear(OpRun):
             if zero_point is not None
             else 0
         )
-        x = x / y_scale
+        if precision:
+            precision_np = tensor_dtype_to_np_dtype(precision)
+            x = x.astype(precision_np) / y_scale.astype(precision_np)
+        else:
+            x = x / y_scale
 
         if tensor_type in _CommonQuantizeLinear.quant_integer_ranges:
             xi = np.rint(x).astype(np.int32)
@@ -199,6 +211,11 @@ class _CommonQuantizeLinear(OpRun):
             i4 = func(xi)
             return (i4,)  # type: ignore[attr-defined]
 
+        if tensor_type == TensorProto.FLOAT4E2M1:
+            x += zero_point
+            f4 = subbyte.float32_to_float4e2m1_unpacked(x)
+            return (f4.astype(float4e2m1),)  # type: ignore[attr-defined]
+
         raise ValueError(
             f"Unexpected type: output_dtype={tensor_type} is not a supported quantized type."
         )
@@ -228,3 +245,24 @@ class QuantizeLinear_21(_CommonQuantizeLinear):
             block_size=block_size,
             output_dtype=output_dtype,
         )  # type: ignore
+
+
+class QuantizeLinear_23(_CommonQuantizeLinear):
+    def _run(
+        self,
+        *args,
+        axis=None,
+        saturate=None,
+        block_size=None,
+        output_dtype=None,
+        precision=None,
+    ):  # type: ignore
+        # args: x, y_scale, zero_point
+        return super()._run(
+            *args,
+            axis=axis,
+            saturate=saturate,
+            block_size=block_size,
+            output_dtype=output_dtype,
+            precision=precision,
+        )  # type: ignore

onnx/reference/ops/op_rms_normalization.py

@@ -0,0 +1,49 @@
+# 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
+
+
+def _rms_normalization(
+    X: np.ndarray,
+    W: np.ndarray,
+    axis: int = -1,
+    epsilon: float = 1e-5,
+) -> np.ndarray:
+    shape = X.shape
+    rank = len(shape)
+    if axis < 0:
+        # If axis = -1 and rank of X is 4,
+        # the axis is changed to -1 + 4 = 3,
+        # which means the last axis.
+        axis = axis + rank
+
+    # This computes RMS for every x_mat's column.
+    x_squared = np.power(X, 2)
+    x_squared_mean = np.mean(
+        x_squared, axis=tuple(range(axis, len(shape))), keepdims=True
+    )
+    # epsilon adjustment to avoid divide-by-zero.
+    rmseps = x_squared_mean + epsilon
+    rms = np.sqrt(rmseps)
+    rms_reciprocal = np.reciprocal(rms)
+
+    y_mat = X * rms_reciprocal
+    # W is linear coefficient.
+    Y = y_mat * W
+
+    return Y.astype(X.dtype)
+
+
+class RMSNormalization(OpRun):
+    def _run(self, X, Scale, axis=None, epsilon=None, stash_type=None):  # type: ignore
+        if stash_type != 1:
+            raise NotImplementedError(
+                f"RMSNormalization not implemented for stash_type={stash_type} != 1."
+            )
+        res = _rms_normalization(X, Scale, axis=axis, epsilon=epsilon)
+        return (res,)  # type: ignore

onnx/reference/ops/op_rotary_embedding.py

@@ -0,0 +1,117 @@
+# 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
+
+
+def rotary_embedding(
+    input: np.ndarray,
+    cos_cache: np.ndarray,
+    sin_cache: np.ndarray,
+    position_ids: np.ndarray | None = None,
+    interleaved=None,
+    rotary_embedding_dim=None,
+    num_heads=None,
+) -> np.ndarray:
+    original_input_shape = input.shape
+    # First ensure input to be processed has shape [batch_size, seq_len, num_heads, head_size]
+    if len(input.shape) == 4:
+        input = np.transpose(input, (0, 2, 1, 3))
+    batch_size = input.shape[0]
+    sequence_length = input.shape[1]
+    if len(input.shape) == 3:
+        hidden_size = input.shape[2]
+        assert num_heads != 0
+        head_size = int(hidden_size / num_heads)
+        new_shape = [batch_size, sequence_length, num_heads, head_size]
+        input = np.reshape(input, new_shape)
+    assert len(input.shape) == 4
+    head_size = input.shape[3]
+
+    # Fully or partially perform rotation on input based on rotary_embedding_dim attribute
+    if rotary_embedding_dim is None or rotary_embedding_dim == 0:
+        # If rotary_embedding_dim not provided, perform full rotation by using head_size
+        rotary_embedding_dim = head_size
+    x_rotate = input[:, :, :, :rotary_embedding_dim]
+    x_not_rotate = input[:, :, :, rotary_embedding_dim:]
+    rotary_embedding_dim_half = int(rotary_embedding_dim / 2)
+
+    # Retrieve sin and cos caches using position ids
+    if position_ids is not None:
+        cos = cos_cache[
+            position_ids
+        ]  # Shape: [batch_size, sequence_length, head_size/2]
+        sin = sin_cache[
+            position_ids
+        ]  # Shape: [batch_size, sequence_length, head_size/2]
+    else:
+        cos = cos_cache
+        sin = sin_cache
+    cos = cos[
+        :, :, :rotary_embedding_dim_half
+    ]  # Shape: [batch_size, sequence_length, rotary_embedding_dim/2]
+    sin = sin[
+        :, :, :rotary_embedding_dim_half
+    ]  # Shape: [batch_size, sequence_length, rotary_embedding_dim/2]
+    cos = np.expand_dims(
+        cos, axis=2
+    )  # Shape: [batch_size, sequence_length, 1, rotary_embedding_dim/2]
+    sin = np.expand_dims(
+        sin, axis=2
+    )  # Shape: [batch_size, sequence_length, 1, rotary_embedding_dim/2]
+
+    # Either divide the input in halves or interleave (based on interleaved attribute)
+    if interleaved:
+        x1 = x_rotate[:, :, :, 0::2]
+        x2 = x_rotate[:, :, :, 1::2]
+    else:
+        x1, x2 = np.split(x_rotate, 2, axis=-1)
+
+    # Calculate real and imaginary values
+    real = (cos * x1) - (sin * x2)
+    imag = (sin * x1) + (cos * x2)
+
+    # Inserted rotated embeddings back to the original input
+    if interleaved:
+        # x_rotate[:, :, :, 0::2] = real
+        # x_rotate[:, :, :, 1::2] = imag
+        real = np.expand_dims(real, axis=-1)
+        imag = np.expand_dims(imag, axis=-1)
+        x_rotate_concat = np.concatenate((real, imag), axis=-1)
+        x_rotate = np.reshape(x_rotate_concat, x_rotate.shape)
+    else:
+        x_rotate = np.concatenate((real, imag), axis=-1)
+    output = np.concatenate((x_rotate, x_not_rotate), axis=-1)
+    if len(original_input_shape) == 3:
+        output = np.reshape(output, original_input_shape)
+    else:
+        output = np.transpose(output, (0, 2, 1, 3))
+    return output
+
+
+class RotaryEmbedding(OpRun):
+    def _run(
+        self,
+        input: np.ndarray,
+        cos_cache: np.ndarray,
+        sin_cache: np.ndarray,
+        position_ids: np.ndarray | None = None,
+        interleaved=None,
+        rotary_embedding_dim=None,
+        num_heads=None,
+    ) -> np.ndarray:
+        return (
+            rotary_embedding(
+                input,
+                cos_cache,
+                sin_cache,
+                position_ids=position_ids,
+                interleaved=interleaved,
+                rotary_embedding_dim=rotary_embedding_dim,
+                num_heads=num_heads,
+            ),
+        )  # type: ignore

onnx/reference/ops/op_scan.py

@@ -127,11 +127,10 @@ class Scan(OpRun):
         results = [[] for _ in scan_names_out]  # type: ignore
 
         for it in range(max_iter):
-            inputs = {}
-            for name, value in zip(state_names_in, states):
-                inputs[name] = value
-            for name, value in zip(scan_names_in, scan_values):
-                inputs[name] = value[it]
+            inputs = dict(zip(state_names_in, states))
+            inputs.update(
+                {name: value[it] for name, value in zip(scan_names_in, scan_values)}
+            )
 
             try:
                 outputs_list = self._run_body(inputs)  # type: ignore

onnx/reference/ops/op_split_to_sequence.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.op_run import OpRun
 
+if TYPE_CHECKING:
+    import numpy as np
+
 
 class SplitToSequence(OpRun):
     def common_run(
@@ -44,8 +47,8 @@ class SplitToSequence(OpRun):
     ) -> tuple[np.ndarray]:
         res = self.common_run(mat, split, axis=axis)
         if split is None and not keepdims:
-            for i in range(len(res)):
-                shape = list(res[i].shape)
+            for i, res_i in enumerate(res):
+                shape = list(res_i.shape)
                 del shape[axis]
-                res[i] = res[i].reshape(tuple(shape))
+                res[i] = res_i.reshape(tuple(shape))
         return (res,)

onnx/reference/ops/op_string_split.py

@@ -14,8 +14,10 @@ def pad_empty_string(
     split_lists: list | np.ndarray, padding_requirement: list | int
 ) -> list:
     if isinstance(split_lists, list):
+        assert isinstance(padding_requirement, int)
         return split_lists + ["" for _ in range(padding_requirement)]
     if isinstance(split_lists, np.ndarray):
+        assert isinstance(padding_requirement, list)
         return list(map(pad_empty_string, split_lists, padding_requirement))
     raise TypeError(f"Invalid array type '{type(split_lists)}'")
 

onnx/reference/ops/op_topk.py

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

onnx/reference/ops/op_unique.py

@@ -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/reference_evaluator.py

@@ -200,15 +200,15 @@ class ReferenceEvaluator:
             implementation a schema may define.
     """
 
-    def __init__(  # type: ignore
+    def __init__(
         self,
         proto: Any,
         opsets: dict[str, int] | None = None,
         functions: list[ReferenceEvaluator | FunctionProto] | None = None,  # type: ignore
         verbose: int = 0,
-        new_ops: list[OpRun] | None = None,
+        new_ops: list[type[OpRun]] | None = None,
         optimized: bool = True,
-    ):
+    ) -> None:
         if optimized:
             if new_ops is None:
                 new_ops = optimized_operators.copy()
@@ -221,7 +221,7 @@ class ReferenceEvaluator:
         self.input_types_ = None
 
         if isinstance(proto, ModelContainer):
-            self.container_ = proto
+            self.container_: ModelContainer | None = proto
             proto = self.container_.model_proto
         else:
             self.container_ = None
@@ -294,7 +294,7 @@ class ReferenceEvaluator:
                 else:
                     raise TypeError(f"Unexpected type {type(f)!r} for a function.")
         self.verbose = verbose
-        self.new_ops_: dict[tuple[str, str], OpRun] = {}
+        self.new_ops_: dict[tuple[str, str], type[OpRun]] = {}
         if new_ops is not None:
             for cl in new_ops:
                 if not hasattr(cl, "op_domain"):
@@ -310,7 +310,7 @@ class ReferenceEvaluator:
                 self.new_ops_[key] = cl
         self._init()
 
-    def retrieve_external_data(self, initializer: TensorProto) -> np.array:
+    def retrieve_external_data(self, initializer: TensorProto) -> np.ndarray:
         """Returns a tensor saved as external."""
         info = ExternalDataInfo(initializer)
         location = info.location
@@ -494,7 +494,7 @@ class ReferenceEvaluator:
                     node.op_type,
                     version,
                     node=node,
-                    input_types=input_types,  # type: ignore[arg-type]
+                    input_types=input_types,
                     expand=expand,
                     evaluator_cls=self.__class__,
                 )
@@ -508,23 +508,17 @@ class ReferenceEvaluator:
         if node.domain == "ai.onnx.preview.training":
             from onnx.reference.ops.aionnx_preview_training import load_op as load_op_pt
 
-            return load_op_pt(
-                node.domain, node.op_type, version, evaluator_cls=self.__class__
-            )
+            return load_op_pt(node.domain, node.op_type, version)
 
         if node.domain == "experimental":
             from onnx.reference.ops.experimental import load_op as load_op_exp
 
-            return load_op_exp(
-                node.domain, node.op_type, version, evaluator_cls=self.__class__
-            )
+            return load_op_exp(node.domain, node.op_type, version)
 
         if node.domain == "ai.onnx.ml":
             from onnx.reference.ops.aionnxml import load_op as load_op_ml
 
-            return load_op_ml(
-                node.domain, node.op_type, version, evaluator_cls=self.__class__
-            )
+            return load_op_ml(node.domain, node.op_type, version)
 
         # It has to be a function.
         if key in self.functions_:

onnx/serialization.py

@@ -11,7 +11,7 @@ __all__ = [
 ]
 
 import typing
-from typing import Any, Collection, Optional, Protocol, TypeVar
+from typing import Any, Optional, Protocol, TypeVar
 
 import google.protobuf.json_format
 import google.protobuf.message
@@ -19,6 +19,9 @@ import google.protobuf.text_format
 
 import onnx
 
+if typing.TYPE_CHECKING:
+    from collections.abc import Collection
+
 _Proto = TypeVar("_Proto", bound=google.protobuf.message.Message)
 # Encoding used for serializing and deserializing text files
 _ENCODING = "utf-8"