onnx 1.15.0__cp311-cp311-win_amd64.whl → 1.16.1__cp311-cp311-win_amd64.whl

onnx/reference/ops/op_hann_window.py

@@ -9,12 +9,9 @@ from onnx.reference.ops._op_common_window import _CommonWindow
 
 
 class HannWindow(_CommonWindow):
-    """
-    Returns
-    :math:`\\omega_n = \\sin^2\\left( \\frac{\\pi n}{N-1} \\right)`
-    where *N* is the window length.
-    See `hann_window
-    <https://pytorch.org/docs/stable/generated/torch.hann_window.html>`_
+    r"""Returns :math:`\\omega_n = \\sin^2\\left( \\frac{\\pi n}{N-1} \\right)` where *N* is the window length.
+
+    See `hann_window <https://pytorch.org/docs/stable/generated/torch.hann_window.html>`_
     """
 
     def _run(self, size, output_datatype=None, periodic=None):  # type: ignore

onnx/reference/ops/op_if.py

@@ -16,10 +16,11 @@ class If(OpRun):
             raise KeyError("run_params must contains key 'opsets'.")
         if "verbose" not in run_params:
             raise KeyError("run_params must contains key 'verbose'.")
+        if "existing_functions" not in self.run_params:
+            raise KeyError("run_params must contains key 'existing_functions'.")
 
     def need_context(self) -> bool:
-        """
-        Tells the runtime if this node needs the context
+        """Tells the runtime if this node needs the context
         (all the results produced so far) as it may silently access
         one of them (operator Loop).
         The default answer is `False`.
@@ -65,4 +66,4 @@ class If(OpRun):
                     f"Output {i!r} (branch={branch!r}, name={names[i]!r}) is None, "
                     f"available inputs={sorted(context)}, initializers={inits}."
                 )
-        return final
+        return self._check_and_fix_outputs(final)

onnx/reference/ops/op_loop.py

@@ -20,8 +20,7 @@ class Loop(OpRun):
         self.K = len(self.body.output_names) - self.N - 1  # type: ignore
 
     def need_context(self) -> bool:
-        """
-        The operator Loop needs to know all results produced
+        """The operator Loop needs to know all results produced
         so far as the loop may silently access one of them.
         Some information are not always referred in the list of inputs
         (kind of static variables).
@@ -34,8 +33,8 @@ class Loop(OpRun):
             args = args[1:]
         else:
             v_initial = None
-        if not hasattr(M, "dtype"):
-            raise TypeError(f"M must be an array or a numpy number not {type(M)}.")
+        if M is not None and not hasattr(M, "dtype"):
+            raise TypeError(f"M must be empty or an array but its type is {type(M)}.")
         body = self.body  # type: ignore
         loop_inputs = body.input_names
         inputs = {name: None for name in loop_inputs}
@@ -53,10 +52,10 @@ class Loop(OpRun):
 
         k_carried_away = [[] for i in range(self.K)]  # type: ignore
         it = 0
-        while cond and it < M:
+        while cond and (M is None or it < M):
             self._log("  -- loop> {%r}", context)
             if len(body.input_names) > 0 and body.input_names[0] is not None:
-                inputs[body.input_names[0]] = np.array(it, dtype=M.dtype)  # type: ignore
+                inputs[body.input_names[0]] = np.array(it, dtype=None if M is None else M.dtype)  # type: ignore
             if len(body.input_names) > 1 and body.input_names[1] is not None:
                 inputs[body.input_names[1]] = cond
             outputs = self._run_body(inputs, attributes=attributes)  # type: ignore
@@ -78,9 +77,8 @@ class Loop(OpRun):
             outputs = [inputs[i] for i in body.input_names[2:]]
         else:
             outputs = outputs[1 : 1 + self.N]
-        outputs.extend(k_carried_away)
+        outputs.extend([np.vstack(x) for x in k_carried_away])
         while len(outputs) < len(self.onnx_node.output):
             outputs.append(np.empty(shape=()))
         res = tuple(outputs)
-
-        return res
+        return self._check_and_fix_outputs(res)

onnx/reference/ops/op_matmul.py

@@ -9,8 +9,7 @@ from onnx.reference.ops._op import OpRunBinaryNum
 
 
 def numpy_matmul(a, b):  # type: ignore
-    """
-    Implements a matmul product. See :func:`np.matmul`.
+    """Implements a matmul product. See :func:`np.matmul`.
     Handles sparse matrices.
     """
     try:

onnx/reference/ops/op_max_pool.py

@@ -87,6 +87,11 @@ class MaxPool(CommonPool):
                         + 1
                     )
                 )
+                need_to_reduce_out_size_in_ceil_mode = (
+                    output_spatial_shape[i] - 1
+                ) * strides[i] >= input_spatial_shape[i] + new_pads[i][0]
+                if need_to_reduce_out_size_in_ceil_mode:
+                    output_spatial_shape[i] -= 1
         else:
             for i in range(len(input_spatial_shape)):
                 output_spatial_shape[i] = int(

onnx/reference/ops/op_optional.py

@@ -8,7 +8,7 @@ from onnx.reference.op_run import OpRun
 
 
 class Optional(OpRun):
-    def _run(self, x=None, type=None):  # type: ignore
+    def _run(self, x=None, type=None):  # type: ignore  # noqa: A002
         if x is not None and type is not None:
             dt = tensor_dtype_to_np_dtype(type)
             if dt != x.dtype:

onnx/reference/ops/op_pool_common.py

@@ -57,8 +57,7 @@ def get_output_shape_explicit_padding(
     dilations: Union[Sequence[int], None] = None,
     ceil_mode: bool = False,
 ) -> Tuple[Sequence[int], Sequence[int]]:
-    """
-    compute output shape according to:
+    """Compute output shape according to:
     https://pytorch.org/docs/stable/generated/torch.nn.MaxPool1d.html?highlight=max+pool#torch.nn.MaxPool1d
     Pads are used to calculate output shape. Use output shape in turn to calculate the actual pads
     that are used to pad the input tensor so that computation in pool() will not cause out of bound error.
@@ -121,8 +120,7 @@ def get_output_shape_auto_pad(
     kernel_spatial_shape: Sequence[int],
     strides_spatial: Sequence[int],
 ) -> Sequence[int]:
-    """
-    https://www.tensorflow.org/api_docs/python/tf/keras/layers/AveragePooling2D
+    """https://www.tensorflow.org/api_docs/python/tf/keras/layers/AveragePooling2D
     output_shape = math.floor((input_shape - 1) / strides) + 1  (SAME)
     output_shape = math.floor((input_shape - pool_size) / strides) + 1 (VALID)
     IMPORTANT: this function assumes ceil_mode is False. In tenforflow, ceil_mode is always False.
@@ -172,8 +170,7 @@ def pool(
     count_include_pad: int = 0,
     p: int = 1,
 ) -> np.ndarray:
-    """
-    this function is used to calculate the pooling result of a padded tensor
+    """This function is used to calculate the pooling result of a padded tensor
     padded: the padded tensor
     x_shape: the shape of the original tensor in [N, C, *spatial_shape]
     kernel: the pooling kernel

onnx/reference/ops/op_qlinear_matmul.py

@@ -22,5 +22,5 @@ class QLinearMatMul(OpRun):
         D = C * (a_scale * b_scale / y_scale)
         if y_zero_point is not None:
             D += y_zero_point
-            return (np.round(D).astype(y_zero_point.dtype),)
-        return (np.round(D).astype(a.dtype),)
+            return (np.rint(D).astype(y_zero_point.dtype),)
+        return (np.rint(D).astype(a.dtype),)

onnx/reference/ops/op_quantize_linear.py

@@ -3,11 +3,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-from typing import Optional, Tuple
+from __future__ import annotations
+
+from typing import ClassVar
 
 import numpy as np
 
-from onnx import TensorProto
+from onnx import TensorProto, subbyte
 from onnx.helper import (
     float32_to_float8e4m3,
     float32_to_float8e5m2,
@@ -19,111 +21,204 @@ from onnx.reference.custom_element_types import (
     float8e4m3fnuz,
     float8e5m2,
     float8e5m2fnuz,
+    int4,
+    uint4,
 )
 from onnx.reference.op_run import OpRun
 
 
+def reshape_input(
+    value: np.ndarray,
+    shape: tuple[int, ...],
+    axis: int | None = None,
+    block_size: int | None = None,
+) -> np.ndarray:
+    """Reshape/Replicate scale/zero-point to be broadcastable to shape.
+
+    Args:
+        value: the array to be reshaped/replicated
+        shape: the rarget shape
+        axis: quantization axis, applicable for per-axis and blocked quantization
+        block_size: size of quantization block, applicable only for blocked quantization
+
+    Returns:
+        value array after reshape/replicate according to quantization mode.
+    """
+    if len(value.shape) == 0:
+        return value
+    if len(value.shape) > 0 and value.size == 1:
+        return value[0]
+    if not block_size:
+        assert len(value.shape) == 1
+        dims = [1] * len(shape)
+        try:
+            dims[axis] = value.size
+            return value.reshape(tuple(dims))
+        except IndexError as e:
+            raise IndexError(
+                f"axis is out of boundary, axis={axis}, "
+                f"value.shape={value.shape}, shape={shape}."
+            ) from e
+
+    if block_size <= 0:
+        raise ValueError("block_size must be a positive integer.")
+
+    # repeat scale to get elementwise scale
+    value = np.repeat(value, repeats=block_size, axis=axis)
+    if (
+        shape[axis] != value.shape[axis]
+    ):  # block_size does not divide x, handle the remainder block
+        value = value.take(indices=range(0, shape[axis]), axis=axis)
+    if value.shape != shape:
+        raise ValueError(
+            "Invalid shapes for Blocked Quantization. Input 2 shape should identical to Input 1 shape, except for one dimension, in which blocking is performed"
+        )
+    assert np.broadcast_shapes(shape, value.shape) == shape
+    return value
+
+
 class _CommonQuantizeLinear(OpRun):
     float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
     float32_to_float8e5m2 = np.vectorize(float32_to_float8e5m2)
+    quant_integer_ranges: ClassVar[dict[TensorProto.DataType, tuple[int]]] = {
+        TensorProto.UINT8: (0, 255),
+        TensorProto.INT8: (-128, 127),
+        TensorProto.UINT16: (0, 65535),
+        TensorProto.INT16: (-32768, 32767),
+    }
+    quant_types = (
+        TensorProto.UINT8,
+        TensorProto.INT8,
+        TensorProto.UINT16,
+        TensorProto.INT16,
+        TensorProto.UINT4,
+        TensorProto.INT4,
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E5M2FNUZ,
+    )
 
     def get_zero_point_type(self, zero_point: np.ndarray) -> int:
+        zero_point_type = None
         if (
             zero_point.dtype == float8e4m3fn
             and zero_point.dtype.descr[0][0] == "e4m3fn"
         ):
-            return TensorProto.FLOAT8E4M3FN
-        if (
+            zero_point_type = TensorProto.FLOAT8E4M3FN
+        elif (
             zero_point.dtype == float8e4m3fnuz
             and zero_point.dtype.descr[0][0] == "e4m3fnuz"
         ):
-            return TensorProto.FLOAT8E4M3FNUZ
-        if zero_point.dtype == float8e5m2 and zero_point.dtype.descr[0][0] == "e5m2":
-            return TensorProto.FLOAT8E5M2
-        if (
+            zero_point_type = TensorProto.FLOAT8E4M3FNUZ
+        elif zero_point.dtype == float8e5m2 and zero_point.dtype.descr[0][0] == "e5m2":
+            zero_point_type = TensorProto.FLOAT8E5M2
+        elif (
             zero_point.dtype == float8e5m2fnuz
             and zero_point.dtype.descr[0][0] == "e5m2fnuz"
         ):
-            return TensorProto.FLOAT8E5M2FNUZ
-        return np_dtype_to_tensor_dtype(zero_point.dtype)
+            zero_point_type = TensorProto.FLOAT8E5M2FNUZ
+        elif zero_point.dtype == uint4 and zero_point.dtype.descr[0][0] == "uint4":
+            zero_point_type = TensorProto.UINT4
+        elif zero_point.dtype == int4 and zero_point.dtype.descr[0][0] == "int4":
+            zero_point_type = TensorProto.INT4
+        else:
+            zero_point_type = np_dtype_to_tensor_dtype(zero_point.dtype)
+        return zero_point_type
 
-    def common_run(  # noqa: PLR0911
+    def _run(
         self,
         x: np.ndarray,
         y_scale: np.ndarray,
-        zero_point: Optional[np.ndarray] = None,
+        zero_point: np.ndarray | None = None,
         axis: int = 1,
         saturate: bool = True,
-    ) -> Tuple[np.ndarray]:
-        if len(y_scale.shape) > 1:
-            raise RuntimeError("Input 2 must be a vector or a number.")
-        if len(y_scale.shape) > 0 and y_scale.size == 1:
-            y_scale = y_scale[0]
-        if len(y_scale.shape) > 0:
-            new_shape = [1 for s in x.shape]
-            new_shape[axis] = len(y_scale)
-            x = x / y_scale.reshape(new_shape)
-        else:
-            x = x / y_scale
-            new_shape = x.shape  # unused
+        block_size: int | None = None,
+        output_dtype: np.dtype | None = None,
+    ) -> tuple[np.ndarray]:
+        y_scale = reshape_input(y_scale, x.shape, axis, block_size)
+
+        # Determine output data type
+        tensor_type = output_dtype
         if zero_point is not None:
-            tensor_type = self.get_zero_point_type(zero_point)
-
-            if tensor_type == TensorProto.UINT8:
-                xi = np.rint(x).astype(np.int32)
-                if len(y_scale.shape) > 0:
-                    xi += zero_point.reshape(new_shape)
-                else:
-                    xi += zero_point
-                dtype = tensor_dtype_to_np_dtype(tensor_type)
-                return (np.clip(xi, 0, 255).astype(dtype),)
-
-            if tensor_type == TensorProto.INT8:
-                xi = np.rint(x).astype(np.int32)
-                if len(y_scale.shape) > 0:
-                    xi += zero_point.reshape(new_shape)
-                else:
-                    xi += zero_point
-                dtype = tensor_dtype_to_np_dtype(tensor_type)
-                return (np.clip(xi, -128, 127).astype(dtype),)
-
-            if tensor_type == TensorProto.FLOAT8E4M3FN:
-                f8 = _CommonQuantizeLinear.float32_to_float8e4m3(x, saturate=saturate)
-                return (f8.astype(float8e4m3fn),)  # type: ignore[attr-defined]
-
-            if tensor_type == TensorProto.FLOAT8E4M3FNUZ:
-                f8 = _CommonQuantizeLinear.float32_to_float8e4m3(
-                    x, uz=True, saturate=saturate
+            zero_point_type = self.get_zero_point_type(zero_point)
+            if output_dtype and output_dtype != zero_point_type:
+                raise ValueError(
+                    f"Mismatched output data-types: output_dtype={output_dtype}, zero_point type={zero_point_type}"
                 )
-                return (f8.astype(float8e4m3fnuz),)  # type: ignore[attr-defined]
+            tensor_type = zero_point_type
+        tensor_type = tensor_type or TensorProto.UINT8
+
+        if tensor_type not in _CommonQuantizeLinear.quant_types:
+            raise ValueError(
+                f"Unexpected type: output_dtype={tensor_type} is not a supported quantized type."
+            )
 
-            if tensor_type == TensorProto.FLOAT8E5M2:
-                f8 = _CommonQuantizeLinear.float32_to_float8e5m2(x, saturate=saturate)
-                return (f8.astype(float8e5m2),)  # type: ignore[attr-defined]
+        # Compute
+        zero_point = (
+            reshape_input(zero_point, x.shape, axis, block_size)
+            if zero_point is not None
+            else 0
+        )
+        x = x / y_scale
+
+        if tensor_type in _CommonQuantizeLinear.quant_integer_ranges:
+            xi = np.rint(x).astype(np.int32)
+            xi += zero_point
+            dtype = tensor_dtype_to_np_dtype(tensor_type)
+            quant_range = _CommonQuantizeLinear.quant_integer_ranges[tensor_type]
+            return (np.clip(xi, quant_range[0], quant_range[1]).astype(dtype),)
+
+        if tensor_type == TensorProto.FLOAT8E4M3FN:
+            f8 = _CommonQuantizeLinear.float32_to_float8e4m3(x, saturate=saturate)
+            return (f8.astype(float8e4m3fn),)  # type: ignore[attr-defined]
+
+        if tensor_type == TensorProto.FLOAT8E4M3FNUZ:
+            f8 = _CommonQuantizeLinear.float32_to_float8e4m3(
+                x, uz=True, saturate=saturate
+            )
+            return (f8.astype(float8e4m3fnuz),)  # type: ignore[attr-defined]
 
-            if tensor_type == TensorProto.FLOAT8E5M2FNUZ:
-                f8 = _CommonQuantizeLinear.float32_to_float8e5m2(
-                    x, fn=True, uz=True, saturate=saturate
-                )
-                return (f8.astype(float8e5m2fnuz),)  # type: ignore[attr-defined]
+        if tensor_type == TensorProto.FLOAT8E5M2:
+            f8 = _CommonQuantizeLinear.float32_to_float8e5m2(x, saturate=saturate)
+            return (f8.astype(float8e5m2),)  # type: ignore[attr-defined]
+
+        if tensor_type == TensorProto.FLOAT8E5M2FNUZ:
+            f8 = _CommonQuantizeLinear.float32_to_float8e5m2(
+                x, fn=True, uz=True, saturate=saturate
+            )
+            return (f8.astype(float8e5m2fnuz),)  # type: ignore[attr-defined]
 
-            raise RuntimeError(
-                f"Unexpected tensor_type for input 2: tensor_type={tensor_type}, "
-                f"zero_point.dtype={zero_point.dtype}."
+        if tensor_type in (TensorProto.UINT4, TensorProto.INT4):
+            xi = np.rint(x).astype(np.int32)
+            xi += zero_point
+            single_func = lambda x: subbyte.float32_to_4bit_unpacked(  # noqa: E731
+                x, signed=(tensor_type == TensorProto.INT4)
             )
+            func = np.vectorize(single_func)
+            i4 = func(xi)
+            return (i4,)  # type: ignore[attr-defined]
 
-        dtype = np.uint8  # type: ignore[assignment]
-        xi = np.rint(x).astype(np.int32)
-        return (np.clip(xi, 0, 255).astype(dtype),)
+        raise ValueError(
+            f"Unexpected type: output_dtype={tensor_type} is not a supported quantized type."
+        )
 
 
 class QuantizeLinear_10(_CommonQuantizeLinear):
-    def _run(self, *args, axis=None):  # type: ignore
-        # args: x, y_scale, zero_point
-        return self.common_run(*args, axis=axis)  # type: ignore
+    def _run(self, x, y_scale, zero_point=None, axis=None):  # type: ignore
+        if len(y_scale.shape) > 1:
+            raise ValueError("Input 2 must be a vector or a number.")
+        return super()._run(x, y_scale, zero_point, axis=axis)  # type: ignore
 
 
 class QuantizeLinear_19(_CommonQuantizeLinear):
-    def _run(self, *args, axis=None, saturate=None):  # type: ignore
+    def _run(self, x, y_scale, zero_point=None, axis=None, saturate=None):  # type: ignore
+        if len(y_scale.shape) > 1:
+            raise ValueError("Input 2 must be a vector or a number.")
+        return super()._run(x, y_scale, zero_point, axis=axis, saturate=saturate)  # type: ignore
+
+
+class QuantizeLinear_21(_CommonQuantizeLinear):
+    def _run(self, *args, axis=None, saturate=None, block_size=None, output_dtype=None):  # type: ignore
         # args: x, y_scale, zero_point
-        return self.common_run(*args, axis=axis, saturate=saturate)  # type: ignore
+        return super()._run(*args, axis=axis, saturate=saturate, block_size=block_size, output_dtype=output_dtype)  # type: ignore

onnx/reference/ops/op_resize.py

@@ -11,8 +11,7 @@ from onnx.reference.op_run import OpRun
 
 
 def _cartesian(arrays: list[np.ndarray], out: np.ndarray | None = None) -> np.ndarray:
-    """
-    From https://stackoverflow.com/a/1235363
+    """From https://stackoverflow.com/a/1235363
     Generate a cartesian product of input arrays.
     Parameters
     ----------
@@ -20,12 +19,14 @@ def _cartesian(arrays: list[np.ndarray], out: np.ndarray | None = None) -> np.nd
         1-D arrays to form the cartesian product of.
     out : ndarray
         Array to place the cartesian product in.
-    Returns
+
+    Returns:
     -------
     out : ndarray
         2-D array of shape (M, len(arrays)) containing cartesian products
         formed of input arrays.
-    Examples
+
+    Examples:
     --------
     >>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
     array([[1, 4, 6],
@@ -41,7 +42,6 @@ def _cartesian(arrays: list[np.ndarray], out: np.ndarray | None = None) -> np.nd
            [3, 5, 6],
            [3, 5, 7]])
     """
-
     arrays = [np.asarray(x) for x in arrays]
     dtype = arrays[0].dtype
 
@@ -127,8 +127,7 @@ def _linear_coeffs_antialias(ratio: float, scale: float) -> np.ndarray:
 
 
 def _get_neighbor_idxes(x: float, n: int, limit: int) -> np.ndarray:
-    """
-    Return the n nearest indexes to x among `[0, limit)`,
+    """Return the n nearest indexes to x among `[0, limit)`,
     prefer the indexes smaller than x.
     As a result, the ratio must be in `(0, 1]`.
 
@@ -142,10 +141,13 @@ def _get_neighbor_idxes(x: float, n: int, limit: int) -> np.ndarray:
         get_neighbor_idxes(4.4, 1, 10) == [4]
         get_neighbor_idxes(4.6, 1, 10) == [5]
 
-    :param x:
-    :param n: the number of the wanted indexes
-    :param limit: the maximum value of index
-    :return: An np.array containing n nearest indexes in ascending order
+    Args:
+        x: float.
+        n: the number of the wanted indexes.
+        limit: the maximum value of index.
+
+    Returns:
+        An np.array containing n nearest indexes in ascending order
     """
     idxes = sorted(range(limit), key=lambda idx: (abs(x - idx), idx))[:n]
     idxes = sorted(idxes)
@@ -153,16 +155,18 @@ def _get_neighbor_idxes(x: float, n: int, limit: int) -> np.ndarray:
 
 
 def _get_neighbor(x: float, n: int, data: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Pad `data` in 'edge' mode, and get n nearest elements in the padded array
-    and their indexes in the original array.
-
-    :param x: center index (in the unpadded coordinate system) of the found nearest elements.
-    :param n: the number of neighbors.
-    :param data: the array
-    :return: A tuple containing the indexes of neighbor elements
-        (the index can be smaller than 0 or higher than len(data))
-        and the value of these elements
+    """Pad `data` in 'edge' mode, and get n nearest elements in the padded array and their indexes in the original array.
+
+    Args:
+        x: Center index (in the unpadded coordinate system) of the found
+            nearest elements.
+        n: The number of neighbors.
+        data: The array.
+
+    Returns:
+        A tuple containing the indexes of neighbor elements (the index
+        can be smaller than 0 or higher than len(data)) and the value of
+        these elements.
     """
     pad_width = np.ceil(n / 2).astype(int)
     padded = np.pad(data, pad_width, mode="edge")

onnx/reference/ops/op_rnn.py

@@ -36,22 +36,30 @@ class CommonRNN(OpRun):
 
         self.f1 = self.choose_act(
             self.activations[0],  # type: ignore
-            self.activation_alpha[0]  # type: ignore
-            if self.activation_alpha is not None and len(self.activation_alpha) > 0  # type: ignore
-            else None,
-            self.activation_beta[0]  # type: ignore
-            if self.activation_beta is not None and len(self.activation_beta) > 0  # type: ignore
-            else None,
+            (
+                self.activation_alpha[0]  # type: ignore
+                if self.activation_alpha is not None and len(self.activation_alpha) > 0  # type: ignore
+                else None
+            ),
+            (
+                self.activation_beta[0]  # type: ignore
+                if self.activation_beta is not None and len(self.activation_beta) > 0  # type: ignore
+                else None
+            ),
         )
         if len(self.activations) > 1:  # type: ignore
             self.f2 = self.choose_act(
                 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
-                else None,
-                self.activation_beta[1]  # type: ignore
-                if self.activation_beta is not None and len(self.activation_beta) > 1  # type: ignore
-                else None,
+                (
+                    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
+                    else None
+                ),
             )
         self.n_outputs = len(onnx_node.output)
 

onnx/reference/ops/op_scan.py

@@ -16,10 +16,12 @@ class Scan(OpRun):
                 f"Parameter 'body' must have a method 'run', type {type(self.body)}."  # type: ignore
             )
         self.input_directions_ = [
-            0
-            if self.scan_input_directions is None  # type: ignore
-            or i >= len(self.scan_input_directions)  # type: ignore
-            else self.scan_input_directions[i]  # type: ignore
+            (
+                0
+                if self.scan_input_directions is None  # type: ignore
+                or i >= len(self.scan_input_directions)  # type: ignore
+                else self.scan_input_directions[i]
+            )  # type: ignore
             for i in range(self.num_scan_inputs)  # type: ignore
         ]
         max_dir_in = max(self.input_directions_)
@@ -28,9 +30,11 @@ class Scan(OpRun):
                 "Scan is not implemented for other output input_direction than 0."
             )
         self.input_axes_ = [
-            0
-            if self.scan_input_axes is None or i >= len(self.scan_input_axes)  # type: ignore
-            else self.scan_input_axes[i]  # type: ignore
+            (
+                0
+                if self.scan_input_axes is None or i >= len(self.scan_input_axes)  # type: ignore
+                else self.scan_input_axes[i]
+            )  # type: ignore
             for i in range(self.num_scan_inputs)  # type: ignore
         ]
         max_axe_in = max(self.input_axes_)
@@ -44,10 +48,12 @@ class Scan(OpRun):
         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
-            else self.scan_output_directions[i]  # type: ignore
+            (
+                0
+                if self.scan_output_directions is None  # type: ignore
+                or i >= len(self.scan_output_directions)  # type: ignore
+                else self.scan_output_directions[i]
+            )  # type: ignore
             for i in range(num_scan_outputs)
         ]
         max_dir_out = max(output_directions)
@@ -56,9 +62,11 @@ class Scan(OpRun):
                 "Scan is not implemented for other output output_direction than 0."
             )
         output_axes = [
-            0
-            if self.scan_output_axes is None or i >= len(self.scan_output_axes)  # type: ignore
-            else self.scan_output_axes[i]  # type: ignore
+            (
+                0
+                if self.scan_output_axes is None or i >= len(self.scan_output_axes)  # type: ignore
+                else self.scan_output_axes[i]
+            )  # type: ignore
             for i in range(num_scan_outputs)
         ]
         max_axe_out = max(output_axes)
@@ -140,4 +148,4 @@ class Scan(OpRun):
         for res in results:
             conc = np.vstack(res)
             states.append(conc)
-        return tuple(states)
+        return self._check_and_fix_outputs(tuple(states))