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

onnx/onnx_operators_ml_pb2.py

@@ -1,11 +1,12 @@
 # -*- coding: utf-8 -*-
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
 # source: onnx/onnx-operators-ml.proto
+# Protobuf Python Version: 4.25.1
 """Generated protocol buffer code."""
-from google.protobuf.internal import builder as _builder
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import descriptor_pool as _descriptor_pool
 from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
 # @@protoc_insertion_point(imports)
 
 _sym_db = _symbol_database.Default()
@@ -16,14 +17,14 @@ from onnx import onnx_ml_pb2 as onnx_dot_onnx__ml__pb2
 
 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1connx/onnx-operators-ml.proto\x12\x04onnx\x1a\x12onnx/onnx-ml.proto\"q\n\rOperatorProto\x12\x0f\n\x07op_type\x18\x01 \x01(\t\x12\x15\n\rsince_version\x18\x02 \x01(\x03\x12$\n\x06status\x18\x03 \x01(\x0e\x32\x14.onnx.OperatorStatus\x12\x12\n\ndoc_string\x18\n \x01(\t\"\xf9\x01\n\x10OperatorSetProto\x12\r\n\x05magic\x18\x01 \x01(\t\x12\x12\n\nir_version\x18\x02 \x01(\x03\x12\x1d\n\x15ir_version_prerelease\x18\x03 \x01(\t\x12\x19\n\x11ir_build_metadata\x18\x07 \x01(\t\x12\x0e\n\x06\x64omain\x18\x04 \x01(\t\x12\x15\n\ropset_version\x18\x05 \x01(\x03\x12\x12\n\ndoc_string\x18\x06 \x01(\t\x12%\n\x08operator\x18\x08 \x03(\x0b\x32\x13.onnx.OperatorProto\x12&\n\tfunctions\x18\t \x03(\x0b\x32\x13.onnx.FunctionProtoB\x02H\x03')
 
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
-_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'onnx.onnx_operators_ml_pb2', globals())
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'onnx.onnx_operators_ml_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
-
-  DESCRIPTOR._options = None
-  DESCRIPTOR._serialized_options = b'H\003'
-  _OPERATORPROTO._serialized_start=58
-  _OPERATORPROTO._serialized_end=171
-  _OPERATORSETPROTO._serialized_start=174
-  _OPERATORSETPROTO._serialized_end=423
+  _globals['DESCRIPTOR']._options = None
+  _globals['DESCRIPTOR']._serialized_options = b'H\003'
+  _globals['_OPERATORPROTO']._serialized_start=58
+  _globals['_OPERATORPROTO']._serialized_end=171
+  _globals['_OPERATORSETPROTO']._serialized_start=174
+  _globals['_OPERATORSETPROTO']._serialized_end=423
 # @@protoc_insertion_point(module_scope)

onnx/onnx_operators_ml_pb2.pyi

@@ -1,67 +1,41 @@
-# @generated by protoc-gen-mypy.py.  Do not edit!
-# mypy: disable-error-code=override
-from google.protobuf.message import (  # type: ignore
-    Message,
-)
+from onnx import onnx_ml_pb2 as _onnx_ml_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
 
-from onnx.onnx_ml_pb2 import (
-    FunctionProto,
-    OperatorStatus,
-)
+DESCRIPTOR: _descriptor.FileDescriptor
 
-from typing import (
-    Iterable,
-    Optional as OptionalType,
-)
+class OperatorProto(_message.Message):
+    __slots__ = ("op_type", "since_version", "status", "doc_string")
+    OP_TYPE_FIELD_NUMBER: _ClassVar[int]
+    SINCE_VERSION_FIELD_NUMBER: _ClassVar[int]
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    DOC_STRING_FIELD_NUMBER: _ClassVar[int]
+    op_type: str
+    since_version: int
+    status: _onnx_ml_pb2.OperatorStatus
+    doc_string: str
+    def __init__(self, op_type: _Optional[str] = ..., since_version: _Optional[int] = ..., status: _Optional[_Union[_onnx_ml_pb2.OperatorStatus, str]] = ..., doc_string: _Optional[str] = ...) -> None: ...
 
-from google.protobuf.internal.containers import (  # type: ignore
-    RepeatedCompositeFieldContainer,
-)
-
-class OperatorProto(Message):
-    op_type = ... # type: str
-    since_version = ... # type: int
-    status = ... # type: OperatorStatus
-    doc_string = ... # type: str
-    
-    def __init__(self,
-        op_type : OptionalType[str] = None,
-        since_version : OptionalType[int] = None,
-        status : OptionalType[OperatorStatus] = None,
-        doc_string : OptionalType[str] = None,
-        ) -> None: ...
-    @classmethod
-    def FromString(cls, s: bytes) -> OperatorProto: ...
-    def MergeFrom(self, other_msg: Message) -> None: ...
-    def CopyFrom(self, other_msg: Message) -> None: ...
-
-class OperatorSetProto(Message):
-    magic = ... # type: str
-    ir_version = ... # type: int
-    ir_version_prerelease = ... # type: str
-    ir_build_metadata = ... # type: str
-    domain = ... # type: str
-    opset_version = ... # type: int
-    doc_string = ... # type: str
-    
-    @property
-    def operator(self) -> RepeatedCompositeFieldContainer[OperatorProto]: ...
-    
-    @property
-    def functions(self) -> RepeatedCompositeFieldContainer[FunctionProto]: ...
-    
-    def __init__(self,
-        magic : OptionalType[str] = None,
-        ir_version : OptionalType[int] = None,
-        ir_version_prerelease : OptionalType[str] = None,
-        ir_build_metadata : OptionalType[str] = None,
-        domain : OptionalType[str] = None,
-        opset_version : OptionalType[int] = None,
-        doc_string : OptionalType[str] = None,
-        operator : OptionalType[Iterable[OperatorProto]] = None,
-        functions : OptionalType[Iterable[FunctionProto]] = None,
-        ) -> None: ...
-    @classmethod
-    def FromString(cls, s: bytes) -> OperatorSetProto: ...
-    def MergeFrom(self, other_msg: Message) -> None: ...
-    def CopyFrom(self, other_msg: Message) -> None: ...
+class OperatorSetProto(_message.Message):
+    __slots__ = ("magic", "ir_version", "ir_version_prerelease", "ir_build_metadata", "domain", "opset_version", "doc_string", "operator", "functions")
+    MAGIC_FIELD_NUMBER: _ClassVar[int]
+    IR_VERSION_FIELD_NUMBER: _ClassVar[int]
+    IR_VERSION_PRERELEASE_FIELD_NUMBER: _ClassVar[int]
+    IR_BUILD_METADATA_FIELD_NUMBER: _ClassVar[int]
+    DOMAIN_FIELD_NUMBER: _ClassVar[int]
+    OPSET_VERSION_FIELD_NUMBER: _ClassVar[int]
+    DOC_STRING_FIELD_NUMBER: _ClassVar[int]
+    OPERATOR_FIELD_NUMBER: _ClassVar[int]
+    FUNCTIONS_FIELD_NUMBER: _ClassVar[int]
+    magic: str
+    ir_version: int
+    ir_version_prerelease: str
+    ir_build_metadata: str
+    domain: str
+    opset_version: int
+    doc_string: str
+    operator: _containers.RepeatedCompositeFieldContainer[OperatorProto]
+    functions: _containers.RepeatedCompositeFieldContainer[_onnx_ml_pb2.FunctionProto]
+    def __init__(self, magic: _Optional[str] = ..., ir_version: _Optional[int] = ..., ir_version_prerelease: _Optional[str] = ..., ir_build_metadata: _Optional[str] = ..., domain: _Optional[str] = ..., opset_version: _Optional[int] = ..., doc_string: _Optional[str] = ..., operator: _Optional[_Iterable[_Union[OperatorProto, _Mapping]]] = ..., functions: _Optional[_Iterable[_Union[_onnx_ml_pb2.FunctionProto, _Mapping]]] = ...) -> None: ...

onnx/parser.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+# ruff: noqa: A005
 from __future__ import annotations
 
 import onnx

onnx/py_utils.h

@@ -15,7 +15,7 @@ template <typename Proto>
 bool ParseProtoFromPyBytes(Proto* proto, const py::bytes& bytes) {
   // Get the buffer from Python bytes object
   char* buffer = nullptr;
-  Py_ssize_t length;
+  Py_ssize_t length = 0;
   PyBytes_AsStringAndSize(bytes.ptr(), &buffer, &length);
 
   return ParseProtoFromBytes(proto, buffer, length);

onnx/reference/custom_element_types.py

@@ -3,15 +3,16 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-import numpy as np
-
 try:
     import ml_dtypes
 except ImportError:
     ml_dtypes = None  # type: ignore[assignment]
 
+from typing import TYPE_CHECKING
+
 from onnx._custom_element_types import (
     bfloat16,
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -20,8 +21,12 @@ from onnx._custom_element_types import (
     uint4,
 )
 
+if TYPE_CHECKING:
+    import numpy as np
+
 _supported_types = [
     (bfloat16, "bfloat16", "bfloat16"),
+    (float4e2m1, "float4e2m1", "float4_e2m1"),
     (float8e4m3fn, "e4m3fn", "float8_e4m3fn"),
     (float8e4m3fnuz, "e4m3fnuz", "float8_e4m3fnuz"),
     (float8e5m2, "e5m2", "float8_e5m2"),
@@ -44,7 +49,8 @@ def convert_from_ml_dtypes(array: np.ndarray) -> np.ndarray:
     if not ml_dtypes:
         return array
     for dtype, _, ml_name in _supported_types:
-        if array.dtype == getattr(ml_dtypes, ml_name):
+        ml_dtype = getattr(ml_dtypes, ml_name, None)
+        if ml_dtype is not None and array.dtype == ml_dtype:
             return array.view(dtype=dtype)
     return array
 

onnx/reference/op_run.py

@@ -4,13 +4,14 @@
 from __future__ import annotations
 
 import abc
-from typing import Any, ClassVar, Iterable
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 
 from onnx import TensorProto
 from onnx._custom_element_types import (
     bfloat16,
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -23,6 +24,9 @@ from onnx.helper import make_node, make_tensor_type_proto, np_dtype_to_tensor_dt
 from onnx.numpy_helper import to_array
 from onnx.onnx_pb import AttributeProto, GraphProto, NodeProto, TypeProto
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 def _split_class_name(name):  # type: ignore
     if "_" in name:
@@ -207,9 +211,9 @@ class OpRun(abc.ABC):
             else:
                 functions = None
             evaluator_cls = self.run_params.get("evaluator_cls", None)
-            assert (
-                evaluator_cls is not None
-            ), f"evaluator_cls must be specified to evaluate att={att}"
+            assert evaluator_cls is not None, (
+                f"evaluator_cls must be specified to evaluate att={att}"
+            )
             return evaluator_cls(
                 att.g,
                 opsets=self.run_params["opsets"],
@@ -254,7 +258,9 @@ class OpRun(abc.ABC):
                 setattr(
                     self,
                     f"_run_{att.name}",
-                    lambda context, value=value, attributes=None: OpRun._evaluate_subgraph(
+                    lambda context,
+                    value=value,
+                    attributes=None: OpRun._evaluate_subgraph(
                         context, value, attributes
                     ),
                 )
@@ -603,12 +609,7 @@ class OpRun(abc.ABC):
 class OpRunExpand(OpRun):
     """Class any operator to avoid must inherit from."""
 
-    def __init__(
-        self,
-        onnx_node: NodeProto,  # noqa: ARG002
-        run_params: dict[str, Any],  # noqa: ARG002
-        impl: Any = None,  # noqa: ARG002
-    ):
+    def __init__(self, *args, **kwargs):  # noqa: ARG002
         raise RuntimeError(
             f"The reference implementation must not use this node ({type(self)})."
         )
@@ -628,7 +629,7 @@ class OpFunction(OpRun):
         run_params: dict[str, Any] | None,
         impl: Any = None,
         attributes: dict[str, Any] | None = None,
-    ):
+    ) -> None:
         if impl is None:
             raise RuntimeError(
                 f"impl cannot be None for node type {onnx_node.op_type!r} "
@@ -650,7 +651,7 @@ class OpFunction(OpRun):
         if len(impl.input_names) != len(inputs):
             raise RuntimeError(
                 f"Mismatch lengths between the number of inputs {len(inputs)} "
-                f"and the expected number of inputs {len(impl.inputs)} "
+                f"and the expected number of inputs {len(impl.input_names)} "
                 f"for node {self.op_type!r} from domain {self.domain!r}."
             )
         feeds = dict(zip(impl.input_names, inputs))
@@ -705,6 +706,8 @@ class OpFunctionContextDependant(OpFunction):
                     ttype = TensorProto.UINT4  # type: ignore[attr-defined]
                 elif t.dtype == int4:
                     ttype = TensorProto.INT4  # type: ignore[attr-defined]
+                elif t.dtype == float4e2m1:
+                    ttype = TensorProto.FLOAT4E2M1  # type: ignore[attr-defined]
                 else:
                     raise
             types.append(make_tensor_type_proto(ttype, t.shape))

onnx/reference/ops/_helpers.py

@@ -12,8 +12,8 @@ from onnx.reference.op_run import OpRun, _split_class_name
 
 def build_registered_operators_any_domain(
     module_context: dict[str, Any],
-) -> dict[str, dict[int | None, OpRun]]:
-    reg_ops: dict[str, dict[int | None, OpRun]] = {}
+) -> dict[str, dict[int | None, type[OpRun]]]:
+    reg_ops: dict[str, dict[int | None, type[OpRun]]] = {}
     for class_name, class_type in module_context.items():
         if class_name.startswith("_") or class_name in {
             "Any",
@@ -31,6 +31,7 @@ def build_registered_operators_any_domain(
             continue
         if isinstance(class_type, type(build_registered_operators_any_domain)):
             continue
+        issub: bool = False
         try:
             issub = issubclass(class_type, OpRun)
         except TypeError as e:

onnx/reference/ops/_op.py

@@ -3,17 +3,19 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
-from onnx.onnx_pb import NodeProto
 from onnx.reference.custom_element_types import (
     convert_from_ml_dtypes,
     convert_to_ml_dtypes,
 )
 from onnx.reference.op_run import OpRun, RuntimeTypeError
 
+if TYPE_CHECKING:
+    from onnx.onnx_pb import NodeProto
+
 
 class OpRunUnary(OpRun):
     """Ancestor to all unary operators in this subfolder.

onnx/reference/ops/_op_list.py

@@ -14,9 +14,7 @@ did not change the implementation.
 from __future__ import annotations
 
 import textwrap
-from typing import Any, Dict, List
-from typing import Optional as TOptional
-from typing import Union
+from typing import Any
 
 from onnx import FunctionProto, NodeProto, TypeProto
 from onnx.defs import get_schema, onnx_opset_version
@@ -26,7 +24,6 @@ from onnx.reference.op_run import (
     OpRun,
     RuntimeContextError,
     RuntimeImplementationError,
-    _split_class_name,
 )
 from onnx.reference.ops._helpers import build_registered_operators_any_domain
 from onnx.reference.ops.op_abs import Abs
@@ -41,6 +38,7 @@ from onnx.reference.ops.op_asin import Asin
 from onnx.reference.ops.op_asinh import Asinh
 from onnx.reference.ops.op_atan import Atan
 from onnx.reference.ops.op_atanh import Atanh
+from onnx.reference.ops.op_attention import Attention
 from onnx.reference.ops.op_attribute_has_value import AttributeHasValue
 from onnx.reference.ops.op_average_pool import (
     AveragePool_1,
@@ -191,8 +189,10 @@ from onnx.reference.ops.op_relu import Relu
 from onnx.reference.ops.op_reshape import Reshape_5, Reshape_14
 from onnx.reference.ops.op_resize import Resize
 from onnx.reference.ops.op_reverse_sequence import ReverseSequence
+from onnx.reference.ops.op_rms_normalization import RMSNormalization
 from onnx.reference.ops.op_rnn import RNN_7, RNN_14
 from onnx.reference.ops.op_roi_align import RoiAlign
+from onnx.reference.ops.op_rotary_embedding import RotaryEmbedding
 from onnx.reference.ops.op_round import Round
 from onnx.reference.ops.op_scan import Scan
 from onnx.reference.ops.op_scatter_elements import ScatterElements
@@ -243,7 +243,7 @@ from onnx.reference.ops.op_where import Where
 from onnx.reference.ops.op_xor import Xor
 
 
-def _build_registered_operators() -> dict[str, dict[int | None, OpRun]]:
+def _build_registered_operators() -> dict[str, dict[int | None, type[OpRun]]]:
     return build_registered_operators_any_domain(globals().copy())
 
 
@@ -279,6 +279,7 @@ def load_op(
     schema = None
     if _registered_operators is None:
         _registered_operators = _build_registered_operators()  # type: ignore[assignment]
+    assert _registered_operators is not None
     if custom is not None:
         return lambda *args: OpFunction(*args, impl=custom)  # type: ignore
     if version is None:
@@ -288,7 +289,7 @@ def load_op(
     if op_type in _registered_operators and not expand:  # type: ignore
         found = True
     else:
-        # maybe the operator can be replacted by a function
+        # maybe the operator can be replaced by a function
         try:
             schema = get_schema(op_type, version, domain)  # type: ignore
         except SchemaError:
@@ -298,9 +299,9 @@ def load_op(
             ) from None
         if schema.has_function:  # type: ignore
             body = schema.function_body  # type: ignore
-            assert (
-                evaluator_cls is not None
-            ), f"evaluator_cls must be specified to implement operator {op_type!r} from domain {domain!r}"
+            assert evaluator_cls is not None, (
+                f"evaluator_cls must be specified to implement operator {op_type!r} from domain {domain!r}"
+            )
             sess = evaluator_cls(body)
             return lambda *args, sess=sess: OpFunction(*args, impl=sess)  # type: ignore
         if schema.has_context_dependent_function:  # type: ignore
@@ -315,9 +316,9 @@ def load_op(
             )
             proto = FunctionProto()
             proto.ParseFromString(body)
-            assert (
-                evaluator_cls is not None
-            ), f"evaluator_cls must be specified to evaluate function {proto.name!r}"
+            assert evaluator_cls is not None, (
+                f"evaluator_cls must be specified to evaluate function {proto.name!r}"
+            )
             sess = evaluator_cls(proto)
             return lambda *args, sess=sess: OpFunction(*args, impl=sess)  # type: ignore
         found = False
@@ -334,7 +335,7 @@ def load_op(
             f"You may either add one or skip the test in "
             f"'reference_evaluator_bakcend_test.py'. Available implementations:\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} "

onnx/reference/ops/aionnx_preview_training/_op_list.py

@@ -4,20 +4,17 @@
 from __future__ import annotations
 
 import textwrap
-from typing import Any, Dict
-from typing import Optional as TOptional
-from typing import Union
+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.aionnx_preview_training._op_run_training import OpRunTraining
 from onnx.reference.ops.aionnx_preview_training.op_adagrad import Adagrad
 from onnx.reference.ops.aionnx_preview_training.op_adam import Adam
 from onnx.reference.ops.aionnx_preview_training.op_momentum import Momentum
 
 
-def _build_registered_operators() -> dict[str, dict[int | None, OpRunTraining]]:
-    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(
@@ -25,7 +22,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.
 
@@ -34,25 +30,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
     if domain != "ai.onnx.preview.training":
         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} "
@@ -84,4 +79,4 @@ def load_op(
     return cl
 
 
-_registered_operators: dict[str, dict[int | None, OpRunTraining]] | None = None
+_registered_operators: dict[str, dict[int | None, type[OpRun]]] | None = None

onnx/reference/ops/aionnxml/_op_list.py

@@ -6,11 +6,9 @@
 from __future__ import annotations
 
 import textwrap
-from typing import Any, Dict
-from typing import Optional as TOptional
-from typing import Union
+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.aionnxml._op_run_aionnxml import OpRunAiOnnxMl
 from onnx.reference.ops.aionnxml.op_array_feature_extractor import ArrayFeatureExtractor
@@ -33,8 +31,8 @@ from onnx.reference.ops.aionnxml.op_tree_ensemble_classifier import (
 from onnx.reference.ops.aionnxml.op_tree_ensemble_regressor import TreeEnsembleRegressor
 
 
-def _build_registered_operators() -> dict[str, dict[int | None, OpRunAiOnnxMl]]:
-    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(
@@ -42,7 +40,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.
 
@@ -51,19 +48,18 @@ 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
     if domain != "ai.onnx.ml":
         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} "
@@ -101,4 +97,4 @@ def load_op(
     return cl
 
 
-_registered_operators: dict[str, dict[int | None, OpRunAiOnnxMl]] | None = None
+_registered_operators: dict[str, dict[int | None, type[OpRun]]] | None = None

onnx/reference/ops/aionnxml/op_dict_vectorizer.py

@@ -11,17 +11,15 @@ from onnx.reference.ops.aionnxml._op_run_aionnxml import OpRunAiOnnxMl
 class DictVectorizer(OpRunAiOnnxMl):
     def _run(self, x, int64_vocabulary=None, string_vocabulary=None):  # type: ignore
         if isinstance(x, (np.ndarray, list)):
-            dict_labels = {}
-            if int64_vocabulary:
-                for i, v in enumerate(int64_vocabulary):
-                    dict_labels[v] = i
-            else:
-                for i, v in enumerate(string_vocabulary):
-                    dict_labels[v] = i
-            if not dict_labels:
-                raise RuntimeError(
-                    "int64_vocabulary and string_vocabulary cannot be both empty."
+            if int64_vocabulary is None and string_vocabulary is None:
+                raise ValueError(
+                    "int64_vocabulary or string_vocabulary must be provided."
                 )
+            if int64_vocabulary is not None:
+                dict_labels = {v: i for i, v in enumerate(int64_vocabulary)}
+            else:
+                assert string_vocabulary is not None
+                dict_labels = {v: i for i, v in enumerate(string_vocabulary)}
 
             values_list = []
             rows_list = []
@@ -40,17 +38,10 @@ class DictVectorizer(OpRunAiOnnxMl):
                 res[r, c] = v
             return (res,)
 
-            # return (
-            #     coo_matrix(
-            #         (values, (rows, cols)), shape=(len(x), len(dict_labels))
-            #     ).todense(),
-            # )
-
         if isinstance(x, dict):
             keys = int64_vocabulary or string_vocabulary
-            result = []
-            for k in keys:
-                result.append(x.get(k, 0))  # noqa: PERF401
+            assert keys is not None
+            result = [x.get(k, 0) for k in keys]
             return (np.array(result),)
 
         raise TypeError(f"x must be iterable not {type(x)}.")

onnx/reference/ops/aionnxml/op_svm_classifier.py

@@ -40,8 +40,7 @@ def multiclass_probability(k, R):
         max_error = 0
         for t in range(k):
             error = np.abs(Qp[t] - pQp)
-            if error > max_error:
-                max_error = error
+            max_error = max(error, max_error)
         if max_error < eps:
             break
 

onnx/reference/ops/aionnxml/op_tree_ensemble_helper.py

@@ -43,7 +43,8 @@ class TreeEnsemble:
 
         self.tree_ids = sorted(set(self.atts.nodes_treeids))  # type: ignore
         self.root_index = {
-            tid: len(self.atts.nodes_treeids) for tid in self.tree_ids  # type: ignore
+            tid: len(self.atts.nodes_treeids)
+            for tid in self.tree_ids  # type: ignore
         }
         for index, tree_id in enumerate(self.atts.nodes_treeids):  # type: ignore
             self.root_index[tree_id] = min(self.root_index[tree_id], index)