onnx 1.17.0__cp312-cp312-macosx_12_0_universal2.whl → 1.19.0__cp312-cp312-macosx_12_0_universal2.whl

onnx/reference/ops/_op_list.py

@@ -1,7 +1,6 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-# pylint: disable=C0415,R0912,R0913,R0914,R0915,W0611,W0603
 """Every class imported in this module defines an implementation of
 an operator of the main domain. Any class name uses `_` to specify a
 version defined in a specific opset. The class name without `_`
@@ -13,10 +12,248 @@ did not change the implementation.
 
 from __future__ import annotations
 
+__all__ = [
+    "load_op",
+    "Abs",
+    "Acos",
+    "Acosh",
+    "Add",
+    "AffineGrid",
+    "And",
+    "ArgMax_1",
+    "ArgMax_12",
+    "ArgMin_1",
+    "ArgMin_12",
+    "Asin",
+    "Asinh",
+    "Atan",
+    "Atanh",
+    "Attention",
+    "AttributeHasValue",
+    "AveragePool_1",
+    "AveragePool_7",
+    "AveragePool_11",
+    "AveragePool_19",
+    "BatchNormalization_6",
+    "BatchNormalization_9",
+    "BatchNormalization_14",
+    "Bernoulli",
+    "BitShift",
+    "BitwiseAnd",
+    "BitwiseNot",
+    "BitwiseOr",
+    "BitwiseXor",
+    "BlackmanWindow",
+    "Cast_1",
+    "Cast_19",
+    "Cast_24",
+    "CastLike_15",
+    "CastLike_19",
+    "Ceil",
+    "Celu",
+    "CenterCropPad",
+    "Clip_6",
+    "Clip_11",
+    "Col2Im",
+    "Compress",
+    "Concat",
+    "ConcatFromSequence",
+    "Constant_1",
+    "Constant_9",
+    "Constant_11",
+    "Constant_12",
+    "ConstantOfShape",
+    "Conv",
+    "ConvInteger",
+    "ConvTranspose",
+    "Cos",
+    "Cosh",
+    "CumSum",
+    "DeformConv",
+    "DepthToSpace",
+    "DequantizeLinear_19",
+    "DequantizeLinear_21",
+    "Det",
+    "DFT_17",
+    "DFT_20",
+    "Div",
+    "Dropout_7",
+    "Dropout_12",
+    "DynamicQuantizeLinear",
+    "Einsum",
+    "Elu",
+    "Equal",
+    "Erf",
+    "Exp",
+    "Expand",
+    "EyeLike",
+    "Flatten",
+    "Floor",
+    "Gather",
+    "GatherElements",
+    "GatherND",
+    "Gemm_6",
+    "Gemm_7",
+    "GlobalAveragePool",
+    "GlobalMaxPool",
+    "Greater",
+    "GreaterOrEqual",
+    "GridSample",
+    "GRU",
+    "HammingWindow",
+    "HannWindow",
+    "HardSigmoid",
+    "Hardmax",
+    "Identity",
+    "If",
+    "ImageDecoder",
+    "InstanceNormalization",
+    "IsInf",
+    "IsNaN",
+    "LayerNormalization",
+    "LeakyRelu",
+    "Less",
+    "LessOrEqual",
+    "Log",
+    "LogSoftmax",
+    "Loop",
+    "LpNormalization",
+    "LpPool",
+    "LRN",
+    "LSTM",
+    "MatMul",
+    "MatMulInteger",
+    "Max",
+    "MaxPool",
+    "MaxUnpool",
+    "Mean",
+    "MelWeightMatrix",
+    "Min",
+    "Mod",
+    "Mul",
+    "Neg",
+    "NegativeLogLikelihoodLoss",
+    "NonMaxSuppression",
+    "NonZero",
+    "Not",
+    "OneHot",
+    "Optional",
+    "OptionalGetElement",
+    "OptionalHasElement",
+    "Or",
+    "Pad_1",
+    "Pad_2",
+    "Pad_11",
+    "Pad_18",
+    "Pow",
+    "PRelu",
+    "QLinearConv",
+    "QLinearMatMul",
+    "QuantizeLinear_10",
+    "QuantizeLinear_19",
+    "QuantizeLinear_21",
+    "RandomNormal",
+    "RandomNormalLike",
+    "RandomUniform",
+    "RandomUniformLike",
+    "Range",
+    "Reciprocal",
+    "ReduceL1_1",
+    "ReduceL1_18",
+    "ReduceL2_1",
+    "ReduceL2_18",
+    "ReduceLogSum_1",
+    "ReduceLogSum_18",
+    "ReduceLogSumExp_1",
+    "ReduceLogSumExp_18",
+    "ReduceMax_1",
+    "ReduceMax_18",
+    "ReduceMean_1",
+    "ReduceMean_18",
+    "ReduceMin_1",
+    "ReduceMin_18",
+    "ReduceProd_1",
+    "ReduceProd_18",
+    "ReduceSum_1",
+    "ReduceSum_13",
+    "ReduceSumSquare_1",
+    "ReduceSumSquare_18",
+    "RegexFullMatch",
+    "Relu",
+    "Reshape_5",
+    "Reshape_14",
+    "Resize",
+    "ReverseSequence",
+    "RMSNormalization",
+    "RNN_7",
+    "RNN_14",
+    "RoiAlign",
+    "RotaryEmbedding",
+    "Round",
+    "Scan",
+    "ScatterElements",
+    "ScatterND",
+    "Selu",
+    "SequenceAt",
+    "SequenceConstruct",
+    "SequenceEmpty",
+    "SequenceErase",
+    "SequenceInsert",
+    "SequenceLength",
+    "SequenceMap",
+    "Shape_1",
+    "Shape_15",
+    "Shrink",
+    "Sigmoid",
+    "Sign",
+    "Sin",
+    "Sinh",
+    "Size",
+    "Slice_1",
+    "Slice_10",
+    "Softmax",
+    "SoftmaxCrossEntropyLoss",
+    "Softplus",
+    "Softsign",
+    "Swish",
+    "SpaceToDepth",
+    "Split_2",
+    "Split_11",
+    "Split_13",
+    "Split_18",
+    "SplitToSequence",
+    "Sqrt",
+    "Squeeze_1",
+    "Squeeze_11",
+    "Squeeze_13",
+    "STFT",
+    "StringConcat",
+    "StringNormalizer",
+    "StringSplit",
+    "Sub",
+    "Sum",
+    "Tan",
+    "Tanh",
+    "TensorScatter",
+    "TfIdfVectorizer",
+    "ThresholdedRelu",
+    "Tile",
+    "TopK_1",
+    "TopK_10",
+    "TopK_11",
+    "Transpose",
+    "Trilu",
+    "Unique",
+    "Unsqueeze_1",
+    "Unsqueeze_11",
+    "Unsqueeze_13",
+    "Upsample",
+    "Where",
+    "Xor",
+]
+
 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 +263,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 +277,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,
@@ -60,7 +297,7 @@ from onnx.reference.ops.op_bitwise_not import BitwiseNot
 from onnx.reference.ops.op_bitwise_or import BitwiseOr
 from onnx.reference.ops.op_bitwise_xor import BitwiseXor
 from onnx.reference.ops.op_blackman_window import BlackmanWindow
-from onnx.reference.ops.op_cast import Cast_1, Cast_19
+from onnx.reference.ops.op_cast import Cast_1, Cast_19, Cast_24
 from onnx.reference.ops.op_cast_like import CastLike_15, CastLike_19
 from onnx.reference.ops.op_ceil import Ceil
 from onnx.reference.ops.op_celu import Celu
@@ -191,8 +428,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
@@ -228,8 +467,10 @@ from onnx.reference.ops.op_string_normalizer import StringNormalizer
 from onnx.reference.ops.op_string_split import StringSplit
 from onnx.reference.ops.op_sub import Sub
 from onnx.reference.ops.op_sum import Sum
+from onnx.reference.ops.op_swish import Swish
 from onnx.reference.ops.op_tan import Tan
 from onnx.reference.ops.op_tanh import Tanh
+from onnx.reference.ops.op_tensor_scatter import TensorScatter
 from onnx.reference.ops.op_tfidf_vectorizer import TfIdfVectorizer
 from onnx.reference.ops.op_thresholded_relu import ThresholdedRelu
 from onnx.reference.ops.op_tile import Tile
@@ -243,7 +484,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())
 
 
@@ -261,13 +502,13 @@ def load_op(
 
     Args:
         domain: domain
-        op_type: oprator type
+        op_type: operator type
         version: requested version
         custom: custom implementation (like a function)
         node: used if no implementation was found and the operator
-            defines a function which is context dependant
+            defines a function which is context dependent
         input_types: used if no implementation was found and the
-            operator defines a function which is context dependant
+            operator defines a function which is context dependent
         expand: use the function implemented in the schema instead of
             its reference implementation
         evaluator_cls: evaluator to use
@@ -279,53 +520,54 @@ 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
+        return lambda *args: OpFunction(*args, impl=custom)
     if version is None:
         version = onnx_opset_version()
     if domain != "":
         raise ValueError(f"Domain must be '' not {domain!r}.")
-    if op_type in _registered_operators and not expand:  # type: ignore
+    if op_type in _registered_operators and not expand:
         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
+            schema = get_schema(op_type, version, domain)
         except SchemaError:
             raise NotImplementedError(
                 f"No registered schema for operator {op_type!r} "
                 f"and domain {domain!r}. Did you recompile the sources after updating the repository?"
             ) 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}"
+        if schema.has_function:
+            body = schema.function_body
+            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
+            return lambda *args, sess=sess: OpFunction(*args, impl=sess)
+        if schema.has_context_dependent_function:
             if node is None or input_types is None:
                 raise RuntimeContextError(
                     f"No registered implementation for operator {op_type!r} "
                     f"and domain {domain!r}, the operator has a context dependent function. "
                     f"but argument node or input_types is not defined (input_types={input_types})."
                 )
-            body = schema.get_context_dependent_function(  # type: ignore
+            body = schema.get_context_dependent_function(
                 node.SerializeToString(), [it.SerializeToString() for it in input_types]
             )
             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
+            return lambda *args, sess=sess: OpFunction(*args, impl=sess)
         found = False
     if not found:
-        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))  # type: ignore
-        has_function = schema.has_function if schema else None  # type: ignore
+        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))
+        has_function = schema.has_function if schema else None
         has_context_dependent_function = (
-            schema.has_context_dependent_function if schema else None  # type: ignore
+            schema.has_context_dependent_function if schema else None
         )
         raise RuntimeImplementationError(
             f"No registered implementation for operator {op_type!r} "
@@ -334,7 +576,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} "
@@ -358,7 +600,7 @@ def load_op(
             )
         cl = impl[best]
     if cl is None:
-        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))  # type: ignore
+        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))
         raise ValueError(
             f"Not registered implementation for operator {op_type!r}, "
             f"domain {domain!r}, and {version!r} in\n{available}"

onnx/reference/ops/aionnx_preview_training/__init__.py

@@ -3,5 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+__all__ = ["load_op", "OpRunTraining"]
+
 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

@@ -3,21 +3,25 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+__all__ = [
+    "load_op",
+    "Adagrad",
+    "Adam",
+    "Momentum",
+]
+
 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,34 +29,32 @@ 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.
 
     Args:
         domain: domain
-        op_type: oprator type
+        op_type: operator type
         version: requested version
         custom: custom implementation (like a function)
-        evaluator_cls: unused
 
     Returns:
         class
     """
     global _registered_operators  # noqa: PLW0603
     if _registered_operators is None:
-        _registered_operators = _build_registered_operators()  # type: ignore[assignment]
+        _registered_operators = _build_registered_operators()
     if custom is not None:
-        return lambda *args: OpFunction(*args, impl=custom)  # type: ignore
+        return lambda *args: OpFunction(*args, impl=custom)
     if domain != "ai.onnx.preview.training":
         raise ValueError(f"Domain must be '' not {domain!r}.")
-    if op_type not in _registered_operators:  # type: ignore
-        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))  # type: ignore
+    if op_type not in _registered_operators:
+        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))
         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} "
@@ -76,7 +78,7 @@ def load_op(
             )
         cl = impl[best]
     if cl is None:
-        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))  # type: ignore
+        available = "\n".join(textwrap.wrap(", ".join(sorted(_registered_operators))))
         raise ValueError(
             f"Not registered implementation for operator {op_type!r}, "
             f"domain {domain!r}, and {version!r} in\n{available}"
@@ -84,4 +86,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/aionnx_preview_training/op_adagrad.py

@@ -8,7 +8,7 @@ import numpy as np
 from onnx.reference.ops.aionnx_preview_training._op_run_training import OpRunTraining
 
 
-def _apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):  # type: ignore
+def _apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):
     # Compute adjusted learning-rate.
     r_ = r / (1 + t * decay_factor)
     # Add gradient of regularization term.
@@ -23,9 +23,9 @@ def _apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):  # t
 
 
 class Adagrad(OpRunTraining):
-    def _run(self, *data, decay_factor=None, epsilon=None, norm_coefficient=None):  # type: ignore
+    def _run(self, *data, decay_factor=None, epsilon=None, norm_coefficient=None):
         if len(data) == 5:
-            return self._run1(  # type: ignore
+            return self._run1(
                 *data,
                 decay_factor=decay_factor,
                 epsilon=epsilon,
@@ -35,7 +35,7 @@ class Adagrad(OpRunTraining):
         xs = []
         hs = []
         for i in range(n):
-            a, b = self._run1(  # type: ignore
+            a, b = self._run1(
                 *data[:2],
                 data[2 + i],
                 data[2 + n + i],
@@ -50,7 +50,7 @@ class Adagrad(OpRunTraining):
 
     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,
@@ -59,6 +59,6 @@ class Adagrad(OpRunTraining):
             h,
             norm_coefficient,
             epsilon,
-            decay_factor,  # type: ignore
+            decay_factor,
         )
         return x_new, h_new

onnx/reference/ops/aionnx_preview_training/op_adam.py

@@ -8,9 +8,9 @@ import numpy as np
 from onnx.reference.ops.aionnx_preview_training._op_run_training import OpRunTraining
 
 
-def _apply_adam(  # type: ignore
+def _apply_adam(
     r, t, x, g, v, h, norm_coefficient, norm_coefficient_post, alpha, beta, epsilon
-):  # type: ignore
+):
     # Add gradient of regularization term.
     g_regularized = norm_coefficient * x + g
     # Update momentum.
@@ -35,7 +35,7 @@ def _apply_adam(  # type: ignore
 
 
 class Adam(OpRunTraining):
-    def _run(  # type: ignore
+    def _run(
         self,
         *data,
         alpha=None,
@@ -45,7 +45,7 @@ class Adam(OpRunTraining):
         norm_coefficient_post=None,
     ):
         if len(data) == 6:
-            return self._run1(  # type: ignore
+            return self._run1(
                 *data,
                 alpha=alpha,
                 beta=beta,
@@ -58,7 +58,7 @@ class Adam(OpRunTraining):
         vs = []
         hs = []
         for i in range(n):
-            a, b, c = self._run1(  # type: ignore
+            a, b, c = self._run1(
                 *data[:2],
                 data[2 + i],
                 data[2 + n + i],
@@ -75,7 +75,7 @@ class Adam(OpRunTraining):
             hs.append(c.astype(np.float32))
         return tuple(xs + vs + hs)
 
-    def _run1(  # type: ignore
+    def _run1(
         self,
         r,
         t,

onnx/reference/ops/aionnx_preview_training/op_momentum.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 from onnx.reference.ops.aionnx_preview_training._op_run_training import OpRunTraining
 
 
-def _apply_momentum(r, t, x, g, v, norm_coefficient, alpha, beta):  # type: ignore
+def _apply_momentum(r, t, x, g, v, norm_coefficient, alpha, beta):
     # Add gradient of regularization term.
     g_regularized = norm_coefficient * x + g
     # Coefficient of gradient should be 1 at the first iteration.
@@ -18,7 +18,7 @@ def _apply_momentum(r, t, x, g, v, norm_coefficient, alpha, beta):  # type: igno
     return x_new, v_new
 
 
-def _apply_nesterov(r, t, x, g, v, norm_coefficient, alpha, beta):  # type: ignore
+def _apply_nesterov(r, t, x, g, v, norm_coefficient, alpha, beta):
     # Add gradient of regularization term.
     g_regularized = norm_coefficient * x + g
     # Coefficient of gradient should be 1 at the first iteration.
@@ -31,10 +31,10 @@ def _apply_nesterov(r, t, x, g, v, norm_coefficient, alpha, beta):  # type: igno
 
 
 class Momentum(OpRunTraining):
-    def _run(self, *data, alpha=None, beta=None, mode=None, norm_coefficient=None):  # type: ignore
+    def _run(self, *data, alpha=None, beta=None, mode=None, norm_coefficient=None):
         if len(data) == 5:
             r, t, x, g, v = data
-            return self._run1(  # type: ignore
+            return self._run1(
                 r,
                 t,
                 x,
@@ -49,7 +49,7 @@ class Momentum(OpRunTraining):
         xs = []
         vs = []
         for i in range(n):
-            a, b = self._run1(  # type: ignore
+            a, b = self._run1(
                 *data[:2],
                 data[2 + i],
                 data[2 + n + i],
@@ -74,7 +74,7 @@ class Momentum(OpRunTraining):
         norm_coefficient=None,
         alpha=None,
         beta=None,
-    ):  # type: ignore
+    ):
         if mode == "standard":
             x_new, v_new = _apply_momentum(r, t, x, g, v, norm_coefficient, alpha, beta)
         else:

onnx/reference/ops/aionnxml/__init__.py

@@ -1,4 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+__all__ = ["load_op"]
+
 from onnx.reference.ops.aionnxml._op_list import load_op