onnx 1.16.2__cp312-cp312-win_amd64.whl → 1.18.0__cp312-cp312-win_amd64.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,8 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+# ruff: noqa: A005
+from __future__ import annotations
 
 import onnx
 import onnx.onnx_cpp2py_export.parser as C  # noqa: N812

onnx/printer.py

@@ -1,14 +1,13 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
-
-from typing import Union
+from __future__ import annotations
 
 import onnx
 import onnx.onnx_cpp2py_export.printer as C  # noqa: N812
 
 
-def to_text(proto: Union[onnx.ModelProto, onnx.FunctionProto, onnx.GraphProto]) -> str:
+def to_text(proto: onnx.ModelProto | onnx.FunctionProto | onnx.GraphProto) -> str:
     if isinstance(proto, onnx.ModelProto):
         return C.model_to_text(proto.SerializeToString())
     if isinstance(proto, onnx.FunctionProto):

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/__init__.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 __all__ = ["ReferenceEvaluator"]
 

onnx/reference/custom_element_types.py

@@ -1,13 +1,84 @@
 # Copyright (c) ONNX Project Contributors
-#
+
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+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,
+    float8e5m2fnuz,
+    int4,
+    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"),
+    (float8e5m2fnuz, "e5m2fnuz", "float8_e5m2fnuz"),
+    (int4, "int4", "int4"),
+    (uint4, "uint4", "uint4"),
+]
+
+
+def convert_from_ml_dtypes(array: np.ndarray) -> np.ndarray:
+    """Detects the type and changes into one of the ONNX
+    defined custom types when ``ml_dtypes`` is installed.
+
+    Args:
+        array: Numpy array with a dtype from ml_dtypes.
+
+    Returns:
+        numpy array
+    """
+    if not ml_dtypes:
+        return array
+    for dtype, _, ml_name in _supported_types:
+        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
+
+
+def convert_to_ml_dtypes(array: np.ndarray) -> np.ndarray:
+    """Detects the type and changes into one of the type
+    defined in ``ml_dtypes`` if installed.
+
+    Args:
+        array: array
+
+    Returns:
+        numpy Numpy array with a dtype from ml_dtypes.
+    """
+    new_dt = None
+
+    for dtype, name, ml_name in _supported_types:
+        if array.dtype == dtype and array.dtype.descr[0][0] == name:
+            assert ml_dtypes, (
+                f"ml_dtypes is not installed and the tensor cannot "
+                f"be converted into ml_dtypes.{array.dtype.descr[0][0]}"
+            )
+            new_dt = getattr(ml_dtypes, ml_name)
+            break
 
-import numpy as np
+    if new_dt:
+        # int4, uint4, the representation uses 1 byte per element,
+        # only onnx storage uses 1 byte for two elements
+        return array.view(dtype=new_dt)
 
-bfloat16 = np.dtype((np.uint16, {"bfloat16": (np.uint16, 0)}))
-float8e4m3fn = np.dtype((np.uint8, {"e4m3fn": (np.uint8, 0)}))
-float8e4m3fnuz = np.dtype((np.uint8, {"e4m3fnuz": (np.uint8, 0)}))
-float8e5m2 = np.dtype((np.uint8, {"e5m2": (np.uint8, 0)}))
-float8e5m2fnuz = np.dtype((np.uint8, {"e5m2fnuz": (np.uint8, 0)}))
-uint4 = np.dtype((np.uint8, {"uint4": (np.uint8, 0)}))
-int4 = np.dtype((np.int8, {"int4": (np.int8, 0)}))
+    return array

onnx/reference/op_run.py

@@ -4,17 +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.defs import get_all_schemas_with_history, get_schema, onnx_opset_version
-from onnx.helper import make_node, make_tensor_type_proto, np_dtype_to_tensor_dtype
-from onnx.numpy_helper import to_array, unpack_int4
-from onnx.onnx_pb import AttributeProto, GraphProto, NodeProto, TypeProto
-from onnx.reference.custom_element_types import (
+from onnx._custom_element_types import (
     bfloat16,
+    float4e2m1,
     float8e4m3fn,
     float8e4m3fnuz,
     float8e5m2,
@@ -22,6 +19,13 @@ from onnx.reference.custom_element_types import (
     int4,
     uint4,
 )
+from onnx.defs import get_all_schemas_with_history, get_schema, onnx_opset_version
+from onnx.helper import make_node, make_tensor_type_proto, np_dtype_to_tensor_dtype
+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
@@ -121,55 +125,7 @@ def to_sparse_tensor(att: AttributeProto) -> SparseTensor:
 
 
 def to_array_extended(tensor: TensorProto) -> np.ndarray:
-    """Similar to :func:`to_array` but deals with non-numpy types bfloat16,
-    float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz, uint4, int4.
-    """
-    elem_type = tensor.data_type
-    if elem_type == TensorProto.BFLOAT16:
-        data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=bfloat16).ravel()
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-
-    if elem_type in (
-        TensorProto.FLOAT8E4M3FN,
-        TensorProto.FLOAT8E4M3FNUZ,
-        TensorProto.FLOAT8E5M2,
-        TensorProto.FLOAT8E5M2FNUZ,
-    ):
-        m = {
-            TensorProto.FLOAT8E4M3FN: float8e4m3fn,
-            TensorProto.FLOAT8E4M3FNUZ: float8e4m3fnuz,
-            TensorProto.FLOAT8E5M2: float8e5m2,
-            TensorProto.FLOAT8E5M2FNUZ: float8e5m2fnuz,
-        }
-
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=m[elem_type]).ravel()  # type: ignore[index]
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-    if elem_type in (TensorProto.UINT4, TensorProto.INT4):
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        m = {TensorProto.INT4: int4, TensorProto.UINT4: uint4}
-        dtype = m[elem_type]  # type: ignore[index]
-        signed = elem_type == TensorProto.INT4
-        y = np.empty(len(data), dtype=dtype).ravel()
-        for i, d in enumerate(data):
-            y[i] = d
-
-        unpacked_data = unpack_int4(y, dims=shape, signed=signed)
-        return unpacked_data.astype(dtype)
+    """Alias for :func:`to_array`."""
     return to_array(tensor)
 
 
@@ -255,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"],
@@ -302,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
                     ),
                 )
@@ -651,14 +609,12 @@ class OpRun(abc.ABC):
 class OpRunExpand(OpRun):
     """Class any operator to avoid must inherit from."""
 
-    def __init__(
-        self, onnx_node: NodeProto, run_params: dict[str, Any], impl: Any = None
-    ):
+    def __init__(self, *args, **kwargs):  # noqa: ARG002
         raise RuntimeError(
             f"The reference implementation must not use this node ({type(self)})."
         )
 
-    def _run(self, *inputs, **kwargs):
+    def _run(self, *inputs, **kwargs):  # noqa: ARG002
         raise RuntimeError(
             f"The reference implementation must not use this node ({type(self)})."
         )
@@ -673,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} "
@@ -695,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))
@@ -735,7 +691,7 @@ class OpFunctionContextDependant(OpFunction):
         for t in inputs:
             try:
                 ttype = np_dtype_to_tensor_dtype(t.dtype)
-            except KeyError as e:
+            except KeyError:
                 if t.dtype == float8e4m3fn:
                     ttype = TensorProto.FLOAT8E4M3FN  # type: ignore[attr-defined]
                 elif t.dtype == float8e4m3fnuz:
@@ -750,8 +706,10 @@ 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 e
+                    raise
             types.append(make_tensor_type_proto(ttype, t.shape))
         cl = self.parent._load_impl(self.onnx_node, types)
         inst = cl(self.onnx_node, self.run_params)

onnx/reference/ops/__init__.py

@@ -3,5 +3,6 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 from onnx.reference.ops._op_list import load_op

onnx/reference/ops/_helpers.py

@@ -3,16 +3,17 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
-from typing import Any, Dict, Union
+from typing import Any
 
 from onnx.reference.op_run import OpRun, _split_class_name
 
 
 def build_registered_operators_any_domain(
-    module_context: Dict[str, Any]
-) -> Dict[str, Dict[Union[int, None], OpRun]]:
-    reg_ops: Dict[str, Dict[Union[int, None], OpRun]] = {}
+    module_context: dict[str, Any],
+) -> 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",
@@ -20,6 +21,7 @@ def build_registered_operators_any_domain(
             "List",
             "TOptional",
             "Union",
+            "annotations",
             "cl",
             "class_name",
             "get_schema",
@@ -29,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

@@ -1,14 +1,21 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
-from typing import Any, Dict
+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.
@@ -22,6 +29,7 @@ class OpRunUnary(OpRun):
         Supports only unary operators.
         """
         self._log("-- begin %s.run(1 input)", self.__class__.__name__)
+        x = convert_to_ml_dtypes(x)
         try:
             res = self._run(x)
         except TypeError as e:
@@ -29,6 +37,7 @@ class OpRunUnary(OpRun):
                 f"Issues with types {', '.join(str(type(_)) for _ in [x])} "
                 f"(unary operator {self.__class__.__name__!r})."
             ) from e
+        res = (convert_from_ml_dtypes(res[0]),)
         self._log("-- done %s.run -> %d outputs", self.__class__.__name__, len(res))
         return self._check_and_fix_outputs(res)
 
@@ -78,6 +87,8 @@ class OpRunBinary(OpRun):
                 f"(operator '{self.__class__.__name__!r}', "
                 f"shapes {x.shape}, {y.shape})."
             )
+        x = convert_to_ml_dtypes(x)
+        y = convert_to_ml_dtypes(y)
         try:
             res = self._run(x, y)
         except (TypeError, ValueError) as e:
@@ -85,6 +96,7 @@ class OpRunBinary(OpRun):
                 f"Issues with types {', '.join(str(type(_)) for _ in [x, y])} "
                 f"(binary operator {self.__class__.__name__!r})."
             ) from e
+        res = (convert_from_ml_dtypes(res[0]),)
         self._log("-- done %s.run -> %d outputs", self.__class__.__name__, len(res))
         return self._check_and_fix_outputs(res)
 
@@ -92,8 +104,6 @@ class OpRunBinary(OpRun):
 class OpRunBinaryComparison(OpRunBinary):
     """Ancestor to all binary operators in this subfolder comparing tensors."""
 
-    pass
-
 
 class OpRunBinaryNum(OpRunBinary):
     """Ancestor to all binary operators in this subfolder.
@@ -119,13 +129,16 @@ class OpRunBinaryNumpy(OpRunBinaryNum):
     """
 
     def __init__(
-        self, numpy_fct: Any, onnx_node: NodeProto, run_params: Dict[str, Any]
+        self, numpy_fct: Any, onnx_node: NodeProto, run_params: dict[str, Any]
     ):
         OpRunBinaryNum.__init__(self, onnx_node, run_params)
         self.numpy_fct = numpy_fct
 
     def _run(self, a, b):  # type: ignore
+        a = convert_to_ml_dtypes(a)
+        b = convert_to_ml_dtypes(b)
         res = (self.numpy_fct(a, b),)
+        res = (convert_from_ml_dtypes(res[0]),)
         return self._check_and_fix_outputs(res)
 
 
@@ -134,7 +147,7 @@ class OpRunReduceNumpy(OpRun):  # type: ignore
     It must have a parameter *axes*.
     """
 
-    def __init__(self, onnx_node: NodeProto, run_params: Dict[str, Any]):
+    def __init__(self, onnx_node: NodeProto, run_params: dict[str, Any]):
         OpRun.__init__(self, onnx_node, run_params)
         if hasattr(self, "axes"):
             if isinstance(self.axes, np.ndarray):  # type: ignore

onnx/reference/ops/_op_common_indices.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np