onnx 1.16.1__cp310-cp310-win32.whl → 1.17.0__cp310-cp310-win32.whl

onnx/reference/custom_element_types.py

@@ -1,13 +1,78 @@
 # Copyright (c) ONNX Project Contributors
-#
+
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 
-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)}))
+try:
+    import ml_dtypes
+except ImportError:
+    ml_dtypes = None  # type: ignore[assignment]
+
+from onnx._custom_element_types import (
+    bfloat16,
+    float8e4m3fn,
+    float8e4m3fnuz,
+    float8e5m2,
+    float8e5m2fnuz,
+    int4,
+    uint4,
+)
+
+_supported_types = [
+    (bfloat16, "bfloat16", "bfloat16"),
+    (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:
+        if array.dtype == getattr(ml_dtypes, ml_name):
+            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
+
+    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)
+
+    return array

onnx/reference/op_run.py

@@ -9,11 +9,7 @@ from typing import Any, ClassVar, Iterable
 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,
     float8e4m3fn,
     float8e4m3fnuz,
@@ -22,6 +18,10 @@ 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
 
 
 def _split_class_name(name):  # type: ignore
@@ -121,55 +121,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)
 
 
@@ -652,13 +604,16 @@ 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
+        self,
+        onnx_node: NodeProto,  # noqa: ARG002
+        run_params: dict[str, Any],  # noqa: ARG002
+        impl: Any = None,  # 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)})."
         )
@@ -735,7 +690,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:
@@ -751,7 +706,7 @@ class OpFunctionContextDependant(OpFunction):
                 elif t.dtype == int4:
                     ttype = TensorProto.INT4  # 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, OpRun]]:
+    reg_ops: dict[str, dict[int | None, 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",

onnx/reference/ops/_op.py

@@ -1,12 +1,17 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
-from typing import Any, Dict
+from typing import 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
 
 
@@ -22,6 +27,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 +35,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 +85,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 +94,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 +102,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 +127,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 +145,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
 

onnx/reference/ops/_op_common_pool.py

@@ -1,10 +1,9 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import itertools
-from typing import Optional, Tuple
 
 import numpy as np
 
@@ -14,11 +13,11 @@ from onnx.reference.ops._op_common_indices import _get_index, _get_indices
 
 def _get_pad_shape(
     auto_pad: str,
-    input_spatial_shape: Tuple[int],
-    kernel_spatial_shape: Tuple[int],
-    strides_spatial: Tuple[int],
-    output_spatial_shape: Tuple[int],
-) -> Tuple[int]:
+    input_spatial_shape: tuple[int],
+    kernel_spatial_shape: tuple[int],
+    strides_spatial: tuple[int],
+    output_spatial_shape: tuple[int],
+) -> tuple[int]:
     pad_shape = [0] * len(input_spatial_shape)
     if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
         for i in range(len(input_spatial_shape)):
@@ -41,10 +40,10 @@ def _get_pad_shape(
 
 def _get_output_shape_no_ceil(
     auto_pad: str,
-    input_spatial_shape: Tuple[int],
-    kernel_spatial_shape: Tuple[int],
-    strides_spatial: Tuple[int],
-) -> Tuple[int]:
+    input_spatial_shape: tuple[int],
+    kernel_spatial_shape: tuple[int],
+    strides_spatial: tuple[int],
+) -> tuple[int]:
     out_shape = [0] * len(input_spatial_shape)
     if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
         for i in range(len(input_spatial_shape)):
@@ -64,12 +63,12 @@ def _get_output_shape_no_ceil(
 
 def _get_output_shape(
     auto_pad: str,
-    input_spatial_shape: Tuple[int],
-    kernel_spatial_shape: Tuple[int],
-    strides_spatial: Tuple[int],
-    pad_shape: Optional[Tuple[int]] = None,
-    ceil_mode: Optional[int] = 0,
-) -> Tuple[int]:
+    input_spatial_shape: tuple[int],
+    kernel_spatial_shape: tuple[int],
+    strides_spatial: tuple[int],
+    pad_shape: tuple[int] | None = None,
+    ceil_mode: int | None = 0,
+) -> tuple[int]:
     if not ceil_mode:
         out_shape = _get_output_shape_no_ceil(
             auto_pad, input_spatial_shape, kernel_spatial_shape, strides_spatial
@@ -119,16 +118,16 @@ def _get_output_shape(
 
 def _pool(
     padded: np.ndarray,
-    x_shape: Tuple[int],
-    kernel_shape: Tuple[int],
-    strides_shape: Tuple[int],
-    out_shape: Tuple[int],
-    pad_shape: Tuple[int],
+    x_shape: tuple[int],
+    kernel_shape: tuple[int],
+    strides_shape: tuple[int],
+    out_shape: tuple[int],
+    pad_shape: tuple[int],
     pooling_type: str,
-    count_include_pad: Optional[int] = 0,
-    ceil_mode: Optional[int] = 0,
+    count_include_pad: int | None = 0,
+    ceil_mode: int | None = 0,
     indices: bool = False,
-    pads: Optional[np.ndarray] = None,
+    pads: np.ndarray | None = None,
 ) -> np.ndarray:
     if pooling_type == "AVG":
         fpool = np.average
@@ -172,7 +171,7 @@ def _pool(
         for i in listi2:
             try:
                 values.append(window[i])
-            except IndexError:
+            except IndexError:  # noqa: PERF203
                 continue
         window_vals = np.array(values)
 
@@ -210,7 +209,7 @@ class CommonPool(OpRun):
         dilations=None,
         kernel_shape=None,
         pads=None,
-        storage_order=None,
+        storage_order=None,  # noqa: ARG002
         strides=None,
     ):
         if pooling_type == "MAX" and dilations is None:
@@ -246,7 +245,12 @@ class CommonPool(OpRun):
         if auto_pad in ("SAME_LOWER", "SAME_UPPER"):
             const = np.nan if count_include_pad == 0 else 0
             out_shape = _get_output_shape(
-                auto_pad, x_shape, kernel_shape, strides, pad_shape, ceil_mode  # type: ignore
+                auto_pad,
+                x_shape,
+                kernel_shape,
+                strides,
+                pad_shape,
+                ceil_mode,  # type: ignore
             )
             pad_shape = _get_pad_shape(  # type: ignore
                 auto_pad, x_shape, kernel_shape, strides, out_shape
@@ -269,7 +273,12 @@ class CommonPool(OpRun):
             )
         else:
             out_shape = _get_output_shape(
-                auto_pad, x_shape, kernel_shape, strides, pad_shape, ceil_mode  # type: ignore
+                auto_pad,
+                x_shape,
+                kernel_shape,
+                strides,
+                pad_shape,
+                ceil_mode,  # type: ignore
             )
 
         n_dims = len(pads) // 2

onnx/reference/ops/_op_common_random.py

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

onnx/reference/ops/_op_common_window.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np
 
@@ -21,6 +21,6 @@ class _CommonWindow(OpRun):
         return ni, N_1
 
     @staticmethod
-    def _end(size, res, output_datatype):  # type: ignore
+    def _end(size, res, output_datatype):  # type: ignore  # noqa: ARG004
         dtype = tensor_dtype_to_np_dtype(output_datatype)
         return (res.astype(dtype),)

onnx/reference/ops/_op_list.py

@@ -10,6 +10,9 @@ with `_`, it means the implementation is valid for every opset.
 The operator may have been updated to support more types but that
 did not change the implementation.
 """
+
+from __future__ import annotations
+
 import textwrap
 from typing import Any, Dict, List
 from typing import Optional as TOptional
@@ -240,19 +243,19 @@ from onnx.reference.ops.op_where import Where
 from onnx.reference.ops.op_xor import Xor
 
 
-def _build_registered_operators() -> Dict[str, Dict[Union[int, None], OpRun]]:
+def _build_registered_operators() -> dict[str, dict[int | None, OpRun]]:
     return build_registered_operators_any_domain(globals().copy())
 
 
 def load_op(
     domain: str,
     op_type: str,
-    version: Union[None, int] = None,
+    version: None | int = None,
     custom: Any = None,
-    node: Union[None, NodeProto] = None,
-    input_types: Union[None, List[TypeProto]] = None,
+    node: None | NodeProto = None,
+    input_types: None | list[TypeProto] = None,
     expand: bool = False,
-    evaluator_cls: TOptional[type] = None,
+    evaluator_cls: type | None = None,
 ) -> Any:
     """Loads the implemented for a specified operator.
 
@@ -363,4 +366,4 @@ def load_op(
     return cl
 
 
-_registered_operators: TOptional[Dict[str, Dict[Union[int, None], OpRun]]] = None
+_registered_operators: dict[str, dict[int | None, OpRun]] | None = None

onnx/reference/ops/aionnx_preview_training/__init__.py

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

onnx/reference/ops/aionnx_preview_training/_op_list.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import textwrap
 from typing import Any, Dict
@@ -16,16 +16,16 @@ 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[Union[int, None], OpRunTraining]]:
+def _build_registered_operators() -> dict[str, dict[int | None, OpRunTraining]]:
     return build_registered_operators_any_domain(globals().copy())  # type: ignore[return-value]
 
 
 def load_op(
     domain: str,
     op_type: str,
-    version: Union[None, int],
+    version: None | int,
     custom: Any = None,
-    evaluator_cls: TOptional[type] = None,
+    evaluator_cls: type | None = None,  # noqa: ARG001
 ) -> Any:
     """Loads the implemented for a specified operator.
 
@@ -84,6 +84,4 @@ def load_op(
     return cl
 
 
-_registered_operators: TOptional[
-    Dict[str, Dict[Union[int, None], OpRunTraining]]
-] = None
+_registered_operators: dict[str, dict[int | None, OpRunTraining]] | None = None

onnx/reference/ops/aionnx_preview_training/_op_run_training.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 from onnx.reference.op_run import OpRun
 

onnx/reference/ops/aionnx_preview_training/op_adagrad.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np
 
@@ -19,7 +19,7 @@ def _apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):  # t
     h_sqrt = np.sqrt(h_new) + epsilon
     # Apply ADAGRAD update rule.
     x_new = x - r_ * g_regularized / h_sqrt
-    return (x_new, h_new)
+    return (x_new.astype(x.dtype), h_new.astype(h.dtype))
 
 
 class Adagrad(OpRunTraining):
@@ -34,7 +34,7 @@ class Adagrad(OpRunTraining):
         n = (len(data) - 2) // 3
         xs = []
         hs = []
-        for i in range(0, n):
+        for i in range(n):
             a, b = self._run1(  # type: ignore
                 *data[:2],
                 data[2 + i],
@@ -48,8 +48,17 @@ class Adagrad(OpRunTraining):
             hs.append(b)
         return tuple(xs + hs)
 
-    def _run1(self, r, t, x, g, h, decay_factor=None, epsilon=None, norm_coefficient=None):  # type: ignore
+    def _run1(
+        self, r, t, x, g, h, decay_factor=None, epsilon=None, norm_coefficient=None
+    ):  # type: ignore
         x_new, h_new = _apply_adagrad(
-            r, t, x, g, h, norm_coefficient, epsilon, decay_factor  # type: ignore
+            r,
+            t,
+            x,
+            g,
+            h,
+            norm_coefficient,
+            epsilon,
+            decay_factor,  # type: ignore
         )
         return x_new, h_new

onnx/reference/ops/aionnx_preview_training/op_adam.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np
 
@@ -57,7 +57,7 @@ class Adam(OpRunTraining):
         xs = []
         vs = []
         hs = []
-        for i in range(0, n):
+        for i in range(n):
             a, b, c = self._run1(  # type: ignore
                 *data[:2],
                 data[2 + i],