onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/reference/ops/op_gather_elements.py

@@ -0,0 +1,45 @@
+# 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 gather_numpy_2(self: np.ndarray, index: np.ndarray) -> np.ndarray:
+    res = []
+    for a, b in zip(self, index):
+        res.append(a[b[0]])
+    return np.array(res, dtype=self.dtype).reshape(index.shape)
+
+
+def gather_numpy(self: np.ndarray, dim: int, index: np.ndarray) -> np.ndarray:
+    idx_xsection_shape = index.shape[:dim] + index.shape[dim + 1 :]
+    self_xsection_shape = self.shape[:dim] + self.shape[dim + 1 :]
+    if idx_xsection_shape != self_xsection_shape:
+        raise ValueError(
+            f"Except for dimension {dim!r}, all dimensions of "
+            f"index and self should be the same size."
+        )
+    data_swaped = np.swapaxes(self, 0, dim)
+    index_swaped = np.swapaxes(index, 0, dim)
+
+    try:
+        gathered = np.choose(index_swaped, data_swaped, mode="wrap")
+    except ValueError:
+        if len(index_swaped.shape) == 2 and len(data_swaped.shape) == 2:
+            return gather_numpy_2(self, index)
+        raise  # pragma: no cover
+
+    return np.swapaxes(gathered, 0, dim)
+
+
+class GatherElements(OpRun):
+    def _run(self, data, indices, axis=None):
+        try:
+            return (gather_numpy(data, axis, indices),)
+        except TypeError:
+            # distribution x86 requires int32.
+            return (gather_numpy(data, axis, indices.astype(int)),)

onnx_diagnostic/reference/ops/op_gather_grad.py

@@ -0,0 +1,12 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+from onnx.reference.ops.op_scatternd import _scatter_nd_impl
+
+
+class GatherGrad(OpRun):
+    op_domain = "com.microsoft"
+
+    def _run(self, shape, indices, updates, reduction=None):
+        data = np.zeros(shape, dtype=updates.dtype)
+        y = _scatter_nd_impl(data, indices, updates, reduction=reduction)
+        return (y,)

onnx_diagnostic/reference/ops/op_memcpy_host.py

@@ -0,0 +1,11 @@
+from onnx.reference.op_run import OpRun
+
+
+class MemcpyFromHost(OpRun):
+    def _run(self, x):
+        return (x,)
+
+
+class MemcpyToHost(OpRun):
+    def _run(self, x):
+        return (x,)

onnx_diagnostic/reference/ops/op_mul_sigmoid.py

@@ -0,0 +1,23 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+def sigmoid(x):  # type: ignore
+    if x > 0:
+        return 1 / (1 + np.exp(-x))
+    return np.exp(x) / (1 + np.exp(x))
+
+
+class MulSigmoid(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def __init__(self, onnx_node, run_params):  # type: ignore
+        OpRun.__init__(self, onnx_node, run_params)
+        self.vf = np.vectorize(sigmoid)
+
+    def _run(self, X):
+        if len(X.shape) == 0:
+            return ((X * sigmoid(X)).astype(X.dtype),)
+        if X.size == 0:
+            return (X,)
+        return ((X * self.vf(X)).astype(X.dtype),)

onnx_diagnostic/reference/ops/op_negxplus1.py

@@ -0,0 +1,8 @@
+from onnx.reference.op_run import OpRun
+
+
+class NegXplus1(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, X):
+        return ((1 - X).astype(X.dtype),)

onnx_diagnostic/reference/ops/op_qlinear_average_pool.py

@@ -0,0 +1,40 @@
+from onnx.reference.op_run import OpRun
+from onnx.reference.ops.op_average_pool import AveragePool_19 as AveragePool
+from onnx.reference.ops.op_dequantize_linear import DequantizeLinear_19 as DequantizeLinear
+from onnx.reference.ops.op_quantize_linear import QuantizeLinear_19 as QuantizeLinear
+
+
+class QLinearAveragePool(OpRun):
+    op_domain = "com.microsoft"
+
+    def _run(
+        self,
+        x,
+        x_scale,
+        x_zero_point,
+        y_scale,
+        y_zero_point,
+        auto_pad=None,
+        ceil_mode=None,
+        channels_last=None,
+        count_include_pad=None,
+        kernel_shape=None,
+        pads=None,
+        strides=None,
+    ):
+        assert channels_last in (
+            None,
+            0,
+        ), f"QLinearAveragePool not implemented if channels_last={channels_last}"
+        dqx = DequantizeLinear.eval(x, x_scale, x_zero_point)
+        y = AveragePool.eval(
+            dqx,
+            auto_pad=auto_pad,
+            ceil_mode=ceil_mode,
+            count_include_pad=count_include_pad,
+            kernel_shape=kernel_shape,
+            pads=pads,
+            strides=strides,
+        )
+        qy = QuantizeLinear.eval(y, y_scale, y_zero_point)
+        return (qy,)

onnx_diagnostic/reference/ops/op_qlinear_conv.py

@@ -0,0 +1,102 @@
+from typing import Tuple
+from onnx.defs import OpSchema
+from onnx.helper import make_attribute
+from onnx.reference.op_run import OpRun
+from onnx.reference.ops.op_conv import Conv
+from onnx.reference.ops.op_dequantize_linear import DequantizeLinear_19 as DequantizeLinear
+from onnx.reference.ops.op_quantize_linear import QuantizeLinear_19 as QuantizeLinear
+
+
+def _switch_dims_nchw_nhwc(dims: Tuple[int, ...], from_nchw_to_nhwc: bool):
+    if len(dims) == 4:
+        if from_nchw_to_nhwc:
+            return (dims[0], *dims[2:], dims[1])
+        return (dims[0], dims[-1], *dims[1:-1])
+    if len(dims) == 3:
+        if from_nchw_to_nhwc:
+            return (*dims[1:], dims[0])
+        return (dims[-1], *dims[:-1])
+    raise NotImplementedError(f"Unable to process shape={dims}")
+
+
+class QLinearConv(OpRun):
+    op_domain = "com.microsoft"
+
+    op_schema = OpSchema(
+        "QLinearConv",
+        "com.microsoft",
+        1,
+        inputs=[
+            OpSchema.FormalParameter("x", "T"),
+            OpSchema.FormalParameter("x_scale", "T"),
+            OpSchema.FormalParameter("x_zero_point", "T1"),
+            OpSchema.FormalParameter("w", "T"),
+            OpSchema.FormalParameter("w_scale", "T"),
+            OpSchema.FormalParameter("w_zero_point", "T2"),
+            OpSchema.FormalParameter("y_scale", "T"),
+            OpSchema.FormalParameter("y_zero_point", "T3"),
+            OpSchema.FormalParameter(
+                "B", "T3", param_option=OpSchema.FormalParameterOption.Optional
+            ),
+        ],
+        outputs=[OpSchema.FormalParameter("y", "T3")],
+        type_constraints=[
+            ("T", ["tensor(float)"], ""),
+            ("T1", ["tensor(int8)", "tensor(uint8)"], ""),
+            ("T2", ["tensor(int8)", "tensor(uint8)"], ""),
+            ("T3", ["tensor(int8)", "tensor(uint8)"], ""),
+        ],
+        attributes=[
+            OpSchema.Attribute("auto_pad", make_attribute("auto_pad", "NOTSET"), ""),
+            OpSchema.Attribute("kernel_shape", OpSchema.AttrType.INTS, "", required=False),
+            OpSchema.Attribute("dilations", OpSchema.AttrType.INTS, "", required=False),
+            OpSchema.Attribute("strides", OpSchema.AttrType.INTS, "", required=False),
+            OpSchema.Attribute("pads", OpSchema.AttrType.INTS, "", required=False),
+            OpSchema.Attribute("group", make_attribute("group", 1), ""),
+            OpSchema.Attribute("channels_last", make_attribute("channels_last", 0), ""),
+        ],
+    )
+
+    def _run(
+        self,
+        x,
+        x_scale,
+        x_zero_point,
+        w,
+        w_scale,
+        w_zero_point,
+        y_scale,
+        y_zero_point,
+        B=None,
+        auto_pad=None,
+        channels_last=None,
+        dilations=None,
+        group=None,
+        kernel_shape=None,
+        pads=None,
+        strides=None,
+    ):
+        dqx = DequantizeLinear.eval(x, x_scale, x_zero_point)
+        dqw = DequantizeLinear.eval(w, w_scale, w_zero_point)
+        if channels_last:
+            dqx = dqx.reshape(_switch_dims_nchw_nhwc(x.shape, False))
+        dqb = (
+            DequantizeLinear.eval(B, x_scale * w_scale, 0).astype(dqx.dtype)
+            if B is not None
+            else None
+        )
+        y = Conv.eval(
+            dqx,
+            dqw,
+            dqb,
+            auto_pad=auto_pad,
+            dilations=dilations,
+            group=group,
+            kernel_shape=kernel_shape,
+            pads=pads,
+            strides=strides,
+        )
+        if channels_last:
+            y = y.reshape(_switch_dims_nchw_nhwc(y.shape, True))
+        qy = QuantizeLinear.eval(y, y_scale, y_zero_point)
+        return (qy,)

onnx_diagnostic/reference/ops/op_quick_gelu.py

@@ -0,0 +1,23 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+def sigmoid(x):  # type: ignore
+    if x > 0:
+        return 1 / (1 + np.exp(-x))
+    return np.exp(x) / (1 + np.exp(x))
+
+
+class QuickGelu(OpRun):
+    op_domain = "com.microsoft"
+
+    def __init__(self, onnx_node, run_params):  # type: ignore
+        OpRun.__init__(self, onnx_node, run_params)
+        self.vf = np.vectorize(sigmoid)
+
+    def _run(self, X, alpha=1.0):
+        if len(X.shape) == 0:
+            return ((X * sigmoid(X * alpha)).astype(X.dtype),)
+        if X.size == 0:
+            return (X,)
+        return ((X * self.vf(X * alpha)).astype(X.dtype),)

onnx_diagnostic/reference/ops/op_replace_zero.py

@@ -0,0 +1,13 @@
+from onnx.reference.op_run import OpRun
+
+
+class ReplaceZero(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, X, by=None, equal=None):
+        x2 = X.copy().flatten()
+        if equal:
+            x2[x2 == 0] = by
+        else:
+            x2[x2 != 0] = by
+        return (x2.reshape(X.shape),)

onnx_diagnostic/reference/ops/op_rotary.py

@@ -0,0 +1,19 @@
+from onnx.reference.op_run import OpRun
+
+
+class Rotary(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, X, splits=None, side=None):
+        assert splits is None or (
+            splits.shape == (2,) and splits[0] == splits[1]
+        ), f"Unexpected split value {splits}"
+        last_dim = X.shape[-1] // 2
+        cp = X.copy()
+        if side == "left":
+            cp[..., :last_dim] = X[..., last_dim:]
+            cp[..., last_dim:] = -X[..., :last_dim]
+        else:
+            cp[..., :last_dim] = -X[..., last_dim:]
+            cp[..., last_dim:] = X[..., :last_dim]
+        return (cp,)

onnx_diagnostic/reference/ops/op_scan.py

@@ -0,0 +1,65 @@
+import numpy as np
+from onnx.reference.ops.op_scan import Scan as _Scan
+
+
+class Scan(_Scan):
+
+    def need_context(self) -> bool:
+        """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`.
+        """
+        return True
+
+    def _run(
+        self,
+        *args,
+        context=None,
+        body=None,
+        num_scan_inputs=None,
+        scan_input_axes=None,
+        scan_input_directions=None,
+        scan_output_axes=None,
+        scan_output_directions=None,
+        attributes=None,
+    ):
+        (
+            num_loop_state_vars,
+            _num_scan_outputs,
+            _output_directions,
+            _max_dir_out,
+            _output_axes,
+            _max_axe_out,
+            state_names_in,
+            state_names_out,
+            scan_names_in,
+            scan_names_out,
+            scan_values,
+            states,
+        ) = 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
+
+        for it in range(max_iter):
+            inputs = context.copy()
+            inputs.update(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
+            except TypeError as e:
+                raise TypeError(
+                    f"Unable to call 'run' for type '{type(self.body)}'."  # type: ignore
+                ) from e
+
+            outputs = dict(zip(self.output_names, outputs_list))
+            states = [outputs[name] for name in state_names_out]
+            for i, name in enumerate(scan_names_out):
+                results[i].append(np.expand_dims(outputs[name], axis=0))
+
+        for res in results:
+            conc = np.vstack(res)
+            states.append(conc)
+        return self._check_and_fix_outputs(tuple(states))

onnx_diagnostic/reference/ops/op_scatter_elements.py

@@ -0,0 +1,107 @@
+import numpy as np
+
+from onnx.reference.op_run import OpRun
+
+
+def scatter_elements(data, indices, updates, axis=0, reduction=None):  # type: ignore
+    if reduction == "add":
+
+        def f(x, y):
+            return x + y
+
+    elif reduction == "min":
+
+        def f(x, y):
+            return np.minimum(x, y)
+
+    elif reduction == "max":
+
+        def f(x, y):
+            return np.maximum(x, y)
+
+    elif reduction == "mul":
+
+        def f(x, y):
+            return x * y
+
+    else:
+
+        def f(x, y):
+            return y
+
+    if axis < 0:
+        axis = data.ndim + axis
+
+    if len(data.shape) == 1 and axis == 0:
+        scattered = np.copy(data)
+        for pos, up in zip(indices, updates):
+            scattered[pos] = f(scattered[pos], up)
+        return scattered
+
+    if len(indices.shape) == 2:
+        scattered = np.copy(data)
+        if axis == 0:
+            for i in range(indices.shape[0]):
+                for j in range(indices.shape[1]):
+                    scattered[indices[i, j], j] = f(scattered[indices[i, j], j], updates[i, j])
+        else:
+            for i in range(indices.shape[0]):
+                for j in range(indices.shape[1]):
+                    scattered[i, indices[i, j]] = f(scattered[i, indices[i, j]], updates[i, j])
+        return scattered
+
+    if len(indices.shape) == 3:
+        scattered = np.copy(data)
+        if axis == 0:
+            for i in range(indices.shape[0]):
+                for j in range(indices.shape[1]):
+                    for k in range(indices.shape[2]):
+                        scattered[indices[i, j, k], j, k] = f(
+                            scattered[indices[i, j, k], j, k], updates[i, j, k]
+                        )
+        elif axis == 1:
+            for i in range(indices.shape[0]):
+                for j in range(indices.shape[1]):
+                    for k in range(indices.shape[2]):
+                        scattered[i, indices[i, j, k], k] = f(
+                            scattered[i, indices[i, j, k], k], updates[i, j, k]
+                        )
+        elif axis == 2:
+            for i in range(indices.shape[0]):
+                for j in range(indices.shape[1]):
+                    for k in range(indices.shape[2]):
+                        scattered[i, j, indices[i, j, k]] = f(
+                            scattered[i, j, indices[i, j, k]], updates[i, j, k]
+                        )
+        return scattered
+
+    if len(indices.shape) == 4:
+        scattered = np.copy(data)
+        if axis == 3:
+            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]):
+                            scattered[a, i, j, indices[a, i, j, k]] = f(
+                                scattered[a, i, j, indices[a, i, j, k]],
+                                updates[a, i, j, k],
+                            )
+            return scattered
+        if axis == 0:
+            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]):
+                            scattered[indices[a, i, j, k], i, j, k] = f(
+                                scattered[indices[a, i, j, k], i, j, k],
+                                updates[a, i, j, k],
+                            )
+            return scattered
+
+    raise RuntimeError(f"Not implemented for indices.shape={indices.shape} and axis={axis}")
+
+
+class ScatterElements(OpRun):
+    def _run(self, data, indices, updates, axis=None, reduction=None):  # type: ignore
+        res = scatter_elements(data, indices, updates, axis=axis, reduction=reduction)
+        return (res,)

onnx_diagnostic/reference/ops/op_scatternd_of_shape.py

@@ -0,0 +1,22 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+from onnx.reference.ops.op_scatternd import _scatter_nd_impl
+
+
+class ScatterNDOfShape(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, shape, indices, updates, reduction=None, strategy=None):
+        data = np.zeros(shape, dtype=updates.dtype)
+        y = _scatter_nd_impl(data, indices, updates, reduction=reduction)
+        return (y,)
+
+
+class MaskedScatterNDOfShape(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, shape, indices, updates, reduction=None, maskedValue=None):
+        data = np.zeros(shape, dtype=updates.dtype)
+        new_updates = np.where(indices == maskedValue, 0, updates)
+        y = _scatter_nd_impl(data, indices, new_updates, reduction=reduction)
+        return (y,)

onnx_diagnostic/reference/ops/op_simplified_layer_normalization.py

@@ -0,0 +1,8 @@
+from onnx.reference.op_run import OpRun
+
+
+class SimplifiedLayerNormalization(OpRun):
+    def _run(self, x, scale, bias=None, axis=None, epsilon=None, stash_type=None):
+        xm = (x**2).mean(axis=axis, keepdims=1) + epsilon
+        xq = xm ** (-0.5)
+        return (x * xq, xq)

onnx_diagnostic/reference/ops/op_skip_layer_normalization.py

@@ -0,0 +1,13 @@
+from onnx.reference.op_run import OpRun
+from onnx.reference.ops.op_layer_normalization import _layer_normalization
+
+
+class SkipLayerNormalization(OpRun):
+    op_domain = "com.microsoft"
+
+    def _run(self, x, skip, gamma=None, beta=None, bias=None, epsilon=None):
+        add = x + skip
+        if bias is not None:
+            add = add + bias
+        res = _layer_normalization(add, gamma, beta, axis=-1, epsilon=epsilon)
+        return (*res, add)

onnx_diagnostic/reference/ops/op_slice.py

@@ -0,0 +1,20 @@
+from onnx.reference.ops.op_slice import SliceCommon
+
+
+class Slice_10(SliceCommon):
+    def __init__(self, onnx_node, run_params):
+        SliceCommon.__init__(self, onnx_node, run_params)
+
+
+class Slice_1(SliceCommon):
+    def __init__(self, onnx_node, run_params):
+        print(onnx_node)
+        SliceCommon.__init__(self, onnx_node, run_params)
+        for f in ["starts", "ends", "steps", "axes"]:
+            if not hasattr(self, f):
+                continue
+            if getattr(self, f) is not None and len(getattr(self, f)) == 0:
+                setattr(self, f, None)
+
+    def _run(self, data, axes=None, ends=None, starts=None):
+        return SliceCommon._run(self, data, starts, ends, axes)

onnx_diagnostic/reference/ops/op_transpose_cast.py

@@ -0,0 +1,16 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class Transpose2DCastFP16(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, X):
+        return (X.T.astype(np.float16),)
+
+
+class Transpose2DCastFP32(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, X):
+        return (X.T.astype(np.float32),)

onnx_diagnostic/reference/ops/op_tri_matrix.py

@@ -0,0 +1,17 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class TriMatrix(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, shape, csts):
+        lower, diag, upper = list(csts)
+        dtype = csts.dtype
+        mat = np.empty(tuple(shape), dtype=dtype)
+        i = np.arange(shape[0], dtype=np.int32).reshape((-1, 1))
+        j = np.arange(shape[1], dtype=np.int32).reshape((1, -1))
+        mat[i > j] = lower
+        mat[i < j] = upper
+        mat[i == j] = diag
+        return (mat,)