onnx 1.17.0__cp311-cp311-win32.whl → 1.19.0__cp311-cp311-win32.whl

onnx/reference/ops/op_reduce_sum.py

@@ -9,7 +9,7 @@ from onnx.reference.ops._op import OpRunReduceNumpy
 
 
 class ReduceSum_1(OpRunReduceNumpy):
-    def _run(self, x, axes=None, keepdims=None):  # type: ignore
+    def _run(self, x, axes=None, keepdims=None):
         axes = tuple(axes) if axes is not None else None
         res = np.sum(x, axis=axes, keepdims=keepdims, dtype=x.dtype)
         if keepdims == 0 and not isinstance(res, np.ndarray):
@@ -19,7 +19,7 @@ class ReduceSum_1(OpRunReduceNumpy):
 
 
 class ReduceSum_13(OpRunReduceNumpy):
-    def _run(self, x, axes=None, keepdims=None, noop_with_empty_axes=None):  # type: ignore
+    def _run(self, x, axes=None, keepdims=None, noop_with_empty_axes=None):
         if (axes is None or axes.shape == (0,)) and noop_with_empty_axes:
             return (x,)
         axes = self.handle_axes(axes)

onnx/reference/ops/op_reduce_sum_square.py

@@ -9,7 +9,7 @@ from onnx.reference.ops._op import OpRunReduceNumpy
 
 
 class ReduceSumSquare_1(OpRunReduceNumpy):
-    def _run(self, data, axes=None, keepdims=None):  # type: ignore
+    def _run(self, data, axes=None, keepdims=None):
         axes = tuple(axes) if axes is not None else None
         res = np.sum(np.square(data), axis=axes, keepdims=keepdims)
         if keepdims == 0 and not isinstance(res, np.ndarray):
@@ -19,12 +19,12 @@ class ReduceSumSquare_1(OpRunReduceNumpy):
 
 
 class ReduceSumSquare_18(OpRunReduceNumpy):
-    def _run(self, data, axes=None, keepdims=1, noop_with_empty_axes=0):  # type: ignore
-        if self.is_axes_empty(axes) and noop_with_empty_axes != 0:  # type: ignore
+    def _run(self, data, axes=None, keepdims=1, noop_with_empty_axes=0):
+        if self.is_axes_empty(axes) and noop_with_empty_axes != 0:
             return (np.square(data),)
 
         axes = self.handle_axes(axes)
-        keepdims = keepdims != 0  # type: ignore
+        keepdims = keepdims != 0
         res = np.sum(np.square(data), axis=axes, keepdims=keepdims)
         if keepdims == 0 and not isinstance(res, np.ndarray):
             # The runtime must return a numpy array of a single float.

onnx/reference/ops/op_regex_full_match.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import re
+
 import numpy as np
 
 from onnx.reference.op_run import OpRun
@@ -12,19 +14,20 @@ _acceptable_str_dtypes = ("U", "O")
 
 class RegexFullMatch(OpRun):
     def _run(self, x, pattern=None):
-        try:
-            import re2
-        except ImportError as e:
-            raise ImportError(
-                "re2 must be installed to use the reference implementation of the RegexFullMatch operator"
-            ) from e
-
-            # As per onnx/mapping.py, object numpy dtype corresponds to TensorProto.STRING
+        # Note: The ONNX specification states that the pattern MUST
+        # follow the re2 syntax. Python's own re package appears to
+        # cover a superset of re2's functionality, albeit definitive
+        # sources are difficult to find. Since the Python bindings to
+        # re2 appear unmaintained, this operator is implemented using
+        # the re module for the time being. This may change if
+        # discrepancies surface in the future.
+
+        # As per onnx/mapping.py, object numpy dtype corresponds to TensorProto.STRING
         if x.dtype.kind not in _acceptable_str_dtypes:
             raise TypeError(f"Input must be string tensor, received dtype {x.dtype}")
         try:
-            regex = re2.compile(pattern)
-        except re2.error as e:
+            regex = re.compile(pattern)
+        except re.error as e:
             raise ValueError(f"Invalid regex pattern {pattern!r}") from e
 
         fullmatch_func = np.vectorize(

onnx/reference/ops/op_relu.py

@@ -9,5 +9,5 @@ from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Relu(OpRunUnaryNum):
-    def _run(self, x):  # type: ignore
+    def _run(self, x):
         return (np.maximum(x, 0).astype(x.dtype),)

onnx/reference/ops/op_reshape.py

@@ -22,7 +22,7 @@ def reshape_reference_implementation(
 
 
 class CommonReshape(OpRun):
-    def _run(self, data, shape):  # type: ignore
+    def _run(self, data, shape):
         return (reshape_reference_implementation(data, shape, 0),)
 
 
@@ -31,7 +31,7 @@ class Reshape_5(CommonReshape):
 
 
 class Reshape_14(CommonReshape):
-    def _run(self, data, shape, allowzero=None):  # type: ignore
+    def _run(self, data, shape, allowzero=None):
         if allowzero is None:
             allowzero = getattr(self, "allowzero", 0) == 1
         else:

onnx/reference/ops/op_resize.py

@@ -7,6 +7,7 @@ from typing import Any, Callable
 
 import numpy as np
 
+import onnx
 from onnx.reference.op_run import OpRun
 
 
@@ -390,7 +391,7 @@ def _interpolate_nd(
 
 
 class Resize(OpRun):
-    def _run(  # type: ignore
+    def _run(
         self,
         X,
         roi,
@@ -431,17 +432,20 @@ class Resize(OpRun):
             raise ValueError(f"Unexpected value {mode!r} for mode.")
 
         if axes is None:
-            output = _interpolate_nd(
-                X,
-                fct,
-                scale_factors=scales,
-                output_size=sizes,
-                roi=roi,
-                keep_aspect_ratio_policy=keep_aspect_ratio_policy,
-                exclude_outside=exclude_outside,
-                coordinate_transformation_mode=coordinate_transformation_mode,  # type: ignore
-                extrapolation_value=extrapolation_value,  # type: ignore
-            ).astype(X.dtype)
+            output = onnx.numpy_helper.saturate_cast(
+                _interpolate_nd(
+                    X,
+                    fct,
+                    scale_factors=scales,
+                    output_size=sizes,
+                    roi=roi,
+                    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
+                    exclude_outside=exclude_outside,
+                    coordinate_transformation_mode=coordinate_transformation_mode,
+                    extrapolation_value=extrapolation_value,
+                ),
+                X.dtype,
+            )
             return (output,)
 
         # axes is not None
@@ -460,14 +464,14 @@ class Resize(OpRun):
                 roi=roi,
                 keep_aspect_ratio_policy=keep_aspect_ratio_policy,
                 exclude_outside=exclude_outside,
-                coordinate_transformation_mode=coordinate_transformation_mode,  # type: ignore
-                extrapolation_value=extrapolation_value,  # type: ignore
-            ).astype(X.dtype)
+                coordinate_transformation_mode=coordinate_transformation_mode,
+                extrapolation_value=extrapolation_value,
+            )
             if res is None:
-                res = np.empty((reshaped.shape[0], *output.shape), dtype=output.dtype)
-            res[i] = output
+                res = np.empty((reshaped.shape[0], *output.shape), dtype=X.dtype)
+            res[i] = onnx.numpy_helper.saturate_cast(output, X.dtype)
 
-        res_reshaped = res.reshape(tuple(X.shape[a] for a in not_axes) + res[0].shape)  # type: ignore
+        res_reshaped = res.reshape(tuple(X.shape[a] for a in not_axes) + res[0].shape)
         new_perm = list(perm)
         for i, a in enumerate(perm):
             new_perm[a] = i

onnx/reference/ops/op_reverse_sequence.py

@@ -7,14 +7,14 @@ from onnx.reference.op_run import OpRun
 
 
 class ReverseSequence(OpRun):
-    def _run(self, data, sequence_lens, batch_axis=None, time_axis=None):  # type: ignore
+    def _run(self, data, sequence_lens, batch_axis=None, time_axis=None):
         index = [slice(0, s) for s in data.shape]
         index_data = [slice(0, s) for s in data.shape]
         result = data.copy()
         for i, sl in enumerate(sequence_lens):
-            index[batch_axis] = i  # type: ignore
+            index[batch_axis] = i
             index[time_axis] = slice(0, sl)
-            index_data[batch_axis] = i  # type: ignore
-            index_data[time_axis] = slice(sl - 1, None, -1)  # type: ignore
+            index_data[batch_axis] = i
+            index_data[time_axis] = slice(sl - 1, None, -1)
             result[tuple(index)] = data[tuple(index_data)]
         return (result,)

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):
+        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,)

onnx/reference/ops/op_rnn.py

@@ -9,73 +9,73 @@ from onnx.reference.op_run import OpRun
 
 
 class CommonRNN(OpRun):
-    def __init__(self, onnx_node, run_params):  # type: ignore
+    def __init__(self, onnx_node, run_params):
         OpRun.__init__(self, onnx_node, run_params)
 
-        if self.direction in ("forward", "reverse"):  # type: ignore
-            self.num_directions = 1  # type: ignore
-        elif self.direction == "bidirectional":  # type: ignore
-            self.num_directions = 2  # type: ignore
+        if self.direction in ("forward", "reverse"):
+            self.num_directions = 1
+        elif self.direction == "bidirectional":
+            self.num_directions = 2
         else:
-            raise RuntimeError(f"Unknown direction {self.direction!r}.")  # type: ignore
+            raise RuntimeError(f"Unknown direction {self.direction!r}.")
 
         if (
-            self.activation_alpha is not None  # type: ignore
-            and len(self.activation_alpha) != self.num_directions  # type: ignore
+            self.activation_alpha is not None
+            and len(self.activation_alpha) != self.num_directions
         ):
             raise RuntimeError(
-                f"activation_alpha must have the same size as num_directions={self.num_directions}."  # type: ignore
+                f"activation_alpha must have the same size as num_directions={self.num_directions}."
             )
         if (
-            self.activation_beta is not None  # type: ignore
-            and len(self.activation_beta) != self.num_directions  # type: ignore
+            self.activation_beta is not None
+            and len(self.activation_beta) != self.num_directions
         ):
             raise RuntimeError(
-                f"activation_beta must have the same size as num_directions={self.num_directions}."  # type: ignore
+                f"activation_beta must have the same size as num_directions={self.num_directions}."
             )
 
         self.f1 = self.choose_act(
-            self.activations[0],  # type: ignore
+            self.activations[0],
             (
-                self.activation_alpha[0]  # type: ignore
-                if self.activation_alpha is not None and len(self.activation_alpha) > 0  # type: ignore
+                self.activation_alpha[0]
+                if self.activation_alpha is not None and len(self.activation_alpha) > 0
                 else None
             ),
             (
-                self.activation_beta[0]  # type: ignore
-                if self.activation_beta is not None and len(self.activation_beta) > 0  # type: ignore
+                self.activation_beta[0]
+                if self.activation_beta is not None and len(self.activation_beta) > 0
                 else None
             ),
         )
-        if len(self.activations) > 1:  # type: ignore
+        if len(self.activations) > 1:
             self.f2 = self.choose_act(
-                self.activations[1],  # type: ignore
+                self.activations[1],
                 (
-                    self.activation_alpha[1]  # type: ignore
+                    self.activation_alpha[1]
                     if self.activation_alpha is not None
-                    and len(self.activation_alpha) > 1  # type: ignore
+                    and len(self.activation_alpha) > 1
                     else None
                 ),
                 (
-                    self.activation_beta[1]  # type: ignore
+                    self.activation_beta[1]
                     if self.activation_beta is not None
-                    and len(self.activation_beta) > 1  # type: ignore
+                    and len(self.activation_beta) > 1
                     else None
                 ),
             )
         self.n_outputs = len(onnx_node.output)
 
-    def choose_act(self, name, alpha, beta):  # type: ignore
+    def choose_act(self, name, alpha, beta):
         if name in ("Tanh", "tanh"):
             return self._f_tanh
         if name in ("Affine", "affine"):
             return lambda x: x * alpha + beta
         raise RuntimeError(f"Unknown activation function {name!r}.")
 
-    def _f_tanh(self, x):  # type: ignore
+    def _f_tanh(self, x):
         return np.tanh(x)
 
-    def _step(self, X, R, B, W, H_0):  # type: ignore
+    def _step(self, X, R, B, W, H_0):
         h_list = []
         H_t = H_0
         for x in np.split(X, X.shape[0], axis=0):
@@ -91,7 +91,7 @@ class CommonRNN(OpRun):
             output = np.expand_dims(concatenated, 1)
         return output, h_list[-1]
 
-    def _run(  # type: ignore
+    def _run(
         self,
         X,
         W,

onnx/reference/ops/op_roi_align.py

@@ -9,7 +9,7 @@ from onnx.reference.op_run import OpRun
 
 
 class PreCalc:
-    def __init__(self, pos1=0, pos2=0, pos3=0, pos4=0, w1=0, w2=0, w3=0, w4=0):  # type: ignore
+    def __init__(self, pos1=0, pos2=0, pos3=0, pos4=0, w1=0, w2=0, w3=0, w4=0):
         self.pos1 = pos1
         self.pos2 = pos2
         self.pos3 = pos3
@@ -25,7 +25,7 @@ class PreCalc:
 
 class RoiAlign(OpRun):
     @staticmethod
-    def pre_calc_for_bilinear_interpolate(  # type: ignore
+    def pre_calc_for_bilinear_interpolate(
         height: int,
         width: int,
         pooled_height: int,
@@ -99,7 +99,7 @@ class RoiAlign(OpRun):
                         w3 = ly * hx
                         w4 = ly * lx
 
-                        # save weights and indeces
+                        # save weights and indices
                         pc = PreCalc()
                         pc.pos1 = y_low * width + x_low
                         pc.pos2 = y_low * width + x_high
@@ -114,7 +114,7 @@ class RoiAlign(OpRun):
                         pre_calc_index += 1
 
     @staticmethod
-    def roi_align_forward(  # type: ignore
+    def roi_align_forward(
         output_shape: tuple[int, int, int, int],
         bottom_data,
         spatial_scale,
@@ -133,7 +133,7 @@ class RoiAlign(OpRun):
         pooled_height = output_shape[2]
         pooled_width = output_shape[3]
 
-        # 100 is a random chosed value, need be tuned
+        # 100 is a random chosen value, need be tuned
         for n in range(n_rois):
             index_n = n * channels * pooled_width * pooled_height
 
@@ -251,7 +251,7 @@ class RoiAlign(OpRun):
 
                         top_data[index] = output_val
 
-    def _run(  # type: ignore
+    def _run(
         self,
         X,
         rois,
@@ -265,12 +265,12 @@ class RoiAlign(OpRun):
     ):
         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
-        sampling_ratio = sampling_ratio or self.sampling_ratio  # type: ignore
-        spatial_scale = spatial_scale or self.spatial_scale  # type: ignore
+        )
+        mode = mode or self.mode
+        output_height = output_height or self.output_height
+        output_width = output_width or self.output_width
+        sampling_ratio = sampling_ratio or self.sampling_ratio
+        spatial_scale = spatial_scale or self.spatial_scale
 
         num_channels = X.shape[1]
         num_rois = batch_indices.shape[0]

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,
+            ),
+        )

onnx/reference/ops/op_round.py

@@ -9,5 +9,5 @@ from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Round(OpRunUnaryNum):
-    def _run(self, x):  # type: ignore
+    def _run(self, x):
         return (np.round(x).astype(x.dtype),)

onnx/reference/ops/op_scan.py

@@ -9,20 +9,20 @@ from onnx.reference.op_run import OpRun
 
 
 class Scan(OpRun):
-    def __init__(self, onnx_node, run_params):  # type: ignore
+    def __init__(self, onnx_node, run_params):
         OpRun.__init__(self, onnx_node, run_params)
-        if not hasattr(self.body, "run"):  # type: ignore
+        if not hasattr(self.body, "run"):
             raise RuntimeError(
-                f"Parameter 'body' must have a method 'run', type {type(self.body)}."  # type: ignore
+                f"Parameter 'body' must have a method 'run', type {type(self.body)}."
             )
         self.input_directions_ = [
             (
                 0
-                if self.scan_input_directions is None  # type: ignore
-                or i >= len(self.scan_input_directions)  # type: ignore
+                if self.scan_input_directions is None
+                or i >= len(self.scan_input_directions)
                 else self.scan_input_directions[i]
-            )  # type: ignore
-            for i in range(self.num_scan_inputs)  # type: ignore
+            )
+            for i in range(self.num_scan_inputs)
         ]
         max_dir_in = max(self.input_directions_)
         if max_dir_in != 0:
@@ -32,28 +32,28 @@ class Scan(OpRun):
         self.input_axes_ = [
             (
                 0
-                if self.scan_input_axes is None or i >= len(self.scan_input_axes)  # type: ignore
+                if self.scan_input_axes is None or i >= len(self.scan_input_axes)
                 else self.scan_input_axes[i]
-            )  # type: ignore
-            for i in range(self.num_scan_inputs)  # type: ignore
+            )
+            for i in range(self.num_scan_inputs)
         ]
         max_axe_in = max(self.input_axes_)
         if max_axe_in != 0:
             raise RuntimeError("Scan is not implemented for other input axes than 0.")
-        self.input_names = self.body.input_names  # type: ignore
-        self.output_names = self.body.output_names  # type: ignore
+        self.input_names = self.body.input_names
+        self.output_names = self.body.output_names
 
-    def _common_run_shape(self, *args):  # type: ignore
-        num_loop_state_vars = len(args) - self.num_scan_inputs  # type: ignore
+    def _common_run_shape(self, *args):
+        num_loop_state_vars = len(args) - self.num_scan_inputs
         num_scan_outputs = len(args) - num_loop_state_vars
 
         output_directions = [
             (
                 0
-                if self.scan_output_directions is None  # type: ignore
-                or i >= len(self.scan_output_directions)  # type: ignore
+                if self.scan_output_directions is None
+                or i >= len(self.scan_output_directions)
                 else self.scan_output_directions[i]
-            )  # type: ignore
+            )
             for i in range(num_scan_outputs)
         ]
         max_dir_out = max(output_directions)
@@ -64,16 +64,16 @@ class Scan(OpRun):
         output_axes = [
             (
                 0
-                if self.scan_output_axes is None or i >= len(self.scan_output_axes)  # type: ignore
+                if self.scan_output_axes is None or i >= len(self.scan_output_axes)
                 else self.scan_output_axes[i]
-            )  # type: ignore
+            )
             for i in range(num_scan_outputs)
         ]
         max_axe_out = max(output_axes)
         if max_axe_out != 0:
             raise RuntimeError("Scan is not implemented for other output axes than 0.")
 
-        state_names_in = self.input_names[: self.num_scan_inputs]  # type: ignore
+        state_names_in = self.input_names[: self.num_scan_inputs]
         state_names_out = self.output_names[: len(state_names_in)]
         scan_names_in = self.input_names[num_loop_state_vars:]
         scan_names_out = self.output_names[num_loop_state_vars:]
@@ -124,20 +124,19 @@ class Scan(OpRun):
         ) = self._common_run_shape(*args)
 
         max_iter = args[num_loop_state_vars].shape[self.input_axes_[0]]
-        results = [[] for _ in scan_names_out]  # type: ignore
+        results = [[] for _ in scan_names_out]
 
         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
+                outputs_list = self._run_body(inputs)
             except TypeError as e:
                 raise TypeError(
-                    f"Unable to call 'run' for type '{type(self.body)}'."  # type: ignore
+                    f"Unable to call 'run' for type '{type(self.body)}'."
                 ) from e
 
             outputs = dict(zip(self.output_names, outputs_list))