onnx 1.16.2__cp311-cp311-win32.whl → 1.18.0__cp311-cp311-win32.whl

onnx/reference/ops/op_regex_full_match.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_relu.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_reshape.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_reverse_sequence.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/op_rms_normalization.py

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

onnx/reference/ops/op_rnn.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np
 
@@ -52,12 +52,14 @@ class CommonRNN(OpRun):
                 self.activations[1],  # type: ignore
                 (
                     self.activation_alpha[1]  # type: ignore
-                    if self.activation_alpha is not None and len(self.activation_alpha) > 1  # type: ignore
+                    if self.activation_alpha is not None
+                    and len(self.activation_alpha) > 1  # type: ignore
                     else None
                 ),
                 (
                     self.activation_beta[1]  # type: ignore
-                    if self.activation_beta is not None and len(self.activation_beta) > 1  # type: ignore
+                    if self.activation_beta is not None
+                    and len(self.activation_beta) > 1  # type: ignore
                     else None
                 ),
             )
@@ -97,11 +99,11 @@ class CommonRNN(OpRun):
         B=None,
         sequence_lens=None,
         initial_h=None,
-        activation_alpha=None,
-        activation_beta=None,
-        activations=None,
-        clip=None,
-        direction=None,
+        activation_alpha=None,  # noqa: ARG002
+        activation_beta=None,  # noqa: ARG002
+        activations=None,  # noqa: ARG002
+        clip=None,  # noqa: ARG002
+        direction=None,  # noqa: ARG002
         hidden_size=None,
         layout=None,
     ):

onnx/reference/ops/op_roi_align.py

@@ -1,9 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
-
-from typing import Tuple
+from __future__ import annotations
 
 import numpy as np
 
@@ -117,7 +115,7 @@ class RoiAlign(OpRun):
 
     @staticmethod
     def roi_align_forward(  # type: ignore
-        output_shape: Tuple[int, int, int, int],
+        output_shape: tuple[int, int, int, int],
         bottom_data,
         spatial_scale,
         height: int,
@@ -265,7 +263,9 @@ class RoiAlign(OpRun):
         sampling_ratio=None,
         spatial_scale=None,
     ):
-        coordinate_transformation_mode = coordinate_transformation_mode or self.coordinate_transformation_mode  # type: ignore
+        coordinate_transformation_mode = (
+            coordinate_transformation_mode or self.coordinate_transformation_mode
+        )  # type: ignore
         mode = mode or self.mode  # type: ignore
         output_height = output_height or self.output_height  # type: ignore
         output_width = output_width or self.output_width  # type: ignore

onnx/reference/ops/op_rotary_embedding.py

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

onnx/reference/ops/op_round.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_scan.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np
 
@@ -99,13 +99,13 @@ class Scan(OpRun):
     def _run(  # type:ignore
         self,
         *args,
-        body=None,
-        num_scan_inputs=None,
-        scan_input_axes=None,
-        scan_input_directions=None,
-        scan_output_axes=None,
-        scan_output_directions=None,
-        attributes=None,
+        body=None,  # noqa: ARG002
+        num_scan_inputs=None,  # noqa: ARG002
+        scan_input_axes=None,  # noqa: ARG002
+        scan_input_directions=None,  # noqa: ARG002
+        scan_output_axes=None,  # noqa: ARG002
+        scan_output_directions=None,  # noqa: ARG002
+        attributes=None,  # noqa: ARG002
     ):
         # TODO: support overridden attributes.
         (
@@ -127,11 +127,10 @@ class Scan(OpRun):
         results = [[] for _ in scan_names_out]  # type: ignore
 
         for it in range(max_iter):
-            inputs = {}
-            for name, value in zip(state_names_in, states):
-                inputs[name] = value
-            for name, value in zip(scan_names_in, scan_values):
-                inputs[name] = value[it]
+            inputs = dict(zip(state_names_in, states))
+            inputs.update(
+                {name: value[it] for name, value in zip(scan_names_in, scan_values)}
+            )
 
             try:
                 outputs_list = self._run_body(inputs)  # type: ignore

onnx/reference/ops/op_scatter_elements.py

@@ -1,7 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import numpy as np
 
@@ -35,7 +35,7 @@ def scatter_elements(data, indices, updates, axis=0, reduction=None):  # type: i
 
     else:
 
-        def f(x, y):
+        def f(x, y):  # noqa: ARG001
             return y
 
     if axis < 0:
@@ -88,55 +88,24 @@ def scatter_elements(data, indices, updates, axis=0, reduction=None):  # type: i
                         )
         return scattered
 
-    idx_xsection_shape = indices.shape[:axis] + indices.shape[axis + 1 :]
-
-    def make_slice(arr, axis, i):  # type: ignore
-        slc = [slice(None)] * arr.ndim
-        slc[axis] = i
-        return slc
-
-    def unpack(packed):  # type: ignore
-        unpacked = packed[0]
-        for i in range(1, len(packed)):
-            unpacked = unpacked, packed[i]
-        return unpacked
-
-    # We use indices and axis parameters to create idx
-    # idx is in a form that can be used as a NumPy advanced
-    # indices for scattering of updates param. in data
-    idx = [
-        [
-            unpack(np.indices(idx_xsection_shape).reshape(indices.ndim - 1, -1)),
-            indices[tuple(make_slice(indices, axis, i))].reshape(1, -1)[0],
-        ]
-        for i in range(indices.shape[axis])
-    ]
-    idx = list(np.concatenate(idx, axis=1))
-    idx.insert(axis, idx.pop())
-
-    # updates_idx is a NumPy advanced indices for indexing
-    # of elements in the updates
-    updates_idx = list(idx)
-    updates_idx.pop(axis)
-    updates_idx.insert(  # type: ignore
-        axis,
-        np.repeat(np.arange(indices.shape[axis]), np.prod(idx_xsection_shape)),  # type: ignore
-    )
+    if len(indices.shape) == 4:
+        scattered = np.copy(data)
+        for a in range(indices.shape[0]):
+            for i in range(indices.shape[1]):
+                for j in range(indices.shape[2]):
+                    for k in range(indices.shape[3]):
+                        index = [a, i, j, k]
+                        index[axis] = indices[a, i, j, k]
+                        tuple_index = tuple(index)
+                        scattered[tuple_index] = f(
+                            scattered[tuple_index],
+                            updates[a, i, j, k],
+                        )
+        return scattered
 
-    scattered = np.copy(data)
-    if reduction == "min":
-        scattered[tuple(idx)] = np.minimum(
-            scattered[tuple(idx)], updates[tuple(updates_idx)]
-        )
-    elif reduction == "max":
-        scattered[tuple(idx)] = np.maximum(
-            scattered[tuple(idx)], updates[tuple(updates_idx)]
-        )
-    elif reduction == "add":
-        scattered[tuple(idx)] += updates[tuple(updates_idx)]
-    else:
-        scattered[tuple(idx)] = updates[tuple(updates_idx)]
-    return scattered
+    raise NotImplementedError(
+        f"ScatterND is not implement for indices.shape={indices.shape} and axis={axis}."
+    )
 
 
 class ScatterElements(OpRun):

onnx/reference/ops/op_scatternd.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_selu.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_sequence_at.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/op_sequence_construct.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/op_sequence_empty.py

@@ -1,11 +1,11 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 from onnx.reference.op_run import OpRun
 
 
 class SequenceEmpty(OpRun):
-    def _run(self, dtype=None):  # type: ignore
+    def _run(self, dtype=None):  # type: ignore  # noqa: ARG002
         return ([],)

onnx/reference/ops/op_sequence_erase.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/op_sequence_insert.py

@@ -1,9 +1,9 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
-
-from typing import Any, List, Optional, Union
+from typing import Any
 
 import numpy as np
 
@@ -11,12 +11,12 @@ from onnx.reference.op_run import OpRun
 
 
 def sequence_insert_reference_implementation(
-    sequence: Union[List[Any], np.ndarray],
+    sequence: list[Any] | np.ndarray,
     tensor: np.ndarray,
-    position: Optional[np.ndarray] = None,
-) -> List[Any]:
+    position: np.ndarray | None = None,
+) -> list[Any]:
     # make a copy of input sequence
-    seq: List[Any] = []
+    seq: list[Any] = []
     if sequence is not None and (
         not isinstance(sequence, np.ndarray) or len(sequence.shape) > 0
     ):

onnx/reference/ops/op_sequence_length.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_sequence_map.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/op_shape.py

@@ -1,9 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
-
-from typing import Optional, Tuple
+from __future__ import annotations
 
 import numpy as np
 
@@ -17,9 +15,7 @@ class Shape_1(OpRun):
 
 class Shape_15(Shape_1):
     @staticmethod
-    def _interval(
-        n: int, start: Optional[int], end: Optional[int]
-    ) -> Optional[Tuple[int, int]]:
+    def _interval(n: int, start: int | None, end: int | None) -> tuple[int, int] | None:
         if start == 0:
             if end is None or np.isnan(end):
                 return None

onnx/reference/ops/op_shrink.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_sigmoid.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_sign.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_sin.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_sinh.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_size.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_slice.py

@@ -1,9 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
-
-
-from typing import Optional
+from __future__ import annotations
 
 import numpy as np
 
@@ -14,8 +12,8 @@ def _slice(
     data: np.ndarray,
     starts: np.ndarray,
     ends: np.ndarray,
-    axes: Optional[np.ndarray] = None,
-    steps: Optional[np.ndarray] = None,
+    axes: np.ndarray | None = None,
+    steps: np.ndarray | None = None,
 ) -> np.ndarray:
     if isinstance(starts, list):
         starts = np.array(starts)

onnx/reference/ops/op_softmax.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_softmax_cross_entropy_loss.py

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