onnx-diagnostic 0.6.1__py3-none-any.whl → 0.6.3__py3-none-any.whl

onnx_diagnostic/reference/torch_evaluator.py

@@ -9,7 +9,7 @@ from . import torch_ops
 
 
 @functools.lru_cache
-def get_kernels() -> Dict[Tuple[str, str, int], type[torch_ops.OpRun]]:
+def get_kernels() -> Dict[Tuple[str, str, int], type[torch_ops.OpRunKernel]]:
     """
     Retrieves all the available kernels class :class:`TorchOnnxEvaluator`
     can use. The full list is the following.
@@ -28,7 +28,7 @@ def get_kernels() -> Dict[Tuple[str, str, int], type[torch_ops.OpRun]]:
     """
     res = {}
     for _k, v in torch_ops.__dict__.items():
-        if isinstance(v, type) and issubclass(v, torch_ops.OpRun) and "_" in v.__name__:
+        if isinstance(v, type) and issubclass(v, torch_ops.OpRunKernel) and "_" in v.__name__:
             name, version = v.__name__.split("_")
             domain = getattr(v, "domain", "")
             res[domain, name, int(version)] = v
@@ -45,6 +45,8 @@ class TorchOnnxEvaluator:
     :param opsets: needed if proto is a graph
     :param functions: known local functions
     :param verbose: verbosity level
+    :param custom_kernels: dictionary of kernels the user can defined to overwrite
+        a specific implementation: ``("", "LayerNormalization"): CustomKernel``
 
     The class holds the following attributes:
 
@@ -98,7 +100,10 @@ class TorchOnnxEvaluator:
         result = sess.run(None, feeds)
         print(string_type(result, with_shape=True, with_min_max=True))
 
-    Adding ``verbose=1`` shows which kernels is executed:
+    With ``verbose=1``, the class prints out every kernel run and
+    and every result deleted along the run.
+    It shows when a result is not needed anymore. In that case,
+    it is deleted to free the memory it takes.
 
     .. runpython::
         :showcode:
@@ -134,8 +139,6 @@ class TorchOnnxEvaluator:
         result = sess.run(None, feeds)
         print(string_type(result, with_shape=True, with_min_max=True))
 
-    It also shows when a result is not needed anymore. In that case,
-    it is deleted to free the memory it takes.
     The runtime can also execute the kernel the onnx model on CUDA.
     It follows the same logic as :class:`onnxruntime.InferenceSession`:
     ``providers=["CUDAExecutionProvider"]``.
@@ -144,6 +147,115 @@ class TorchOnnxEvaluator:
     identified as a shape in CPU. Some bugs may remain as torch
     raises an exception when devices are expected to be the same.
     The runtime was validated with model :epkg:`arnir0/Tiny-LLM`.
+    Next example shows how to replace a kernel with a different
+    one based on :epkg:`onnxruntime`.
+
+    .. runpython::
+        :showcode:
+
+        import numpy as np
+        import onnx
+        import onnx.helper as oh
+        import onnxruntime
+        import torch
+        from onnx_diagnostic.helpers import string_type
+        from onnx_diagnostic.helpers.torch_helper import onnx_dtype_to_torch_dtype
+        from onnx_diagnostic.reference import TorchOnnxEvaluator
+        from onnx_diagnostic.reference.torch_ops import OpRunKernel, OpRunTensor
+
+        TFLOAT16 = onnx.TensorProto.FLOAT16
+
+        class LayerNormalizationOrt(OpRunKernel):
+            "LayerNormalization based on onnxruntime"
+
+            def __init__(self, node: onnx.NodeProto, version=None, verbose=0):
+                super().__init__(node, version, verbose=verbose)
+                self.axis = self.get_attribute_int(node, "axis", -1)
+                self.epsilon = self.get_attribute_float(node, "epsilon", 1e-5)
+                self.stash_type = onnx_dtype_to_torch_dtype(
+                    self.get_attribute_int(node, "stash_type", onnx.TensorProto.FLOAT)
+                )
+                self.compute_std = len(node.output) > 1
+                assert not self.compute_std, "The keren only computes the first output."
+                layer_model = oh.make_model(
+                    oh.make_graph(
+                        [
+                            oh.make_node(
+                                "LayerNormalization",
+                                ["X", "W", "B"],
+                                ["Z"],
+                                axis=-1,
+                                epsilon=9.999999974752427e-7,
+                            )
+                        ],
+                        "dummy",
+                        [
+                            oh.make_tensor_value_info("X", TFLOAT16, ["b", "c", "d"]),
+                            oh.make_tensor_value_info("W", TFLOAT16, ["d"]),
+                            oh.make_tensor_value_info("B", TFLOAT16, ["d"]),
+                        ],
+                        [oh.make_tensor_value_info("Z", TFLOAT16, ["b", "c", "d"])],
+                    ),
+                    ir_version=9,
+                    opset_imports=[oh.make_opsetid("", 17)],
+                )
+                self.ort_sess = onnxruntime.InferenceSession(
+                    layer_model.SerializeToString(), providers=["CUDAExecutionProvider"]
+                )
+
+            def run(self, x, scale, bias=None):
+                print(f"-- running {self.__class__.__name__}")
+                feeds = dict(X=x, W=scale)
+                if bias is not None:
+                    feeds["B"] = bias
+                feeds = {k: v.tensor.detach().cpu().numpy() for k, v in feeds.items()}
+                got = self.ort_sess.run(None, feeds)[0]
+                return OpRunTensor(torch.from_numpy(got).to(x.dtype).to(x.device))
+
+        # This kernel is tested on this model.
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "LayerNormalization",
+                        ["X", "W", "B"],
+                        ["ln"],
+                        axis=-1,
+                        epsilon=9.999999974752427e-7,
+                    ),
+                    oh.make_node(
+                        "Add", ["ln", "W"], ["Z"], axis=-1, epsilon=9.999999974752427e-7
+                    ),
+                ],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT16, ["b", "c", "d"]),
+                    oh.make_tensor_value_info("W", TFLOAT16, ["d"]),
+                    oh.make_tensor_value_info("B", TFLOAT16, ["d"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT16, ["b", "c", "d"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 17)],
+        )
+
+        torch_sess = TorchOnnxEvaluator(
+            model,
+            custom_kernels={("", "LayerNormalization"): LayerNormalizationOrt},
+            verbose=1,
+        )
+        feeds = dict(
+            zip(
+                torch_sess.input_names,
+                [
+                    torch.rand(3, 4, 5, dtype=torch.float16),
+                    torch.abs(torch.rand(5, dtype=torch.float16)),
+                    torch.rand(5, dtype=torch.float16),
+                ],
+            )
+        )
+        res = torch_sess.run(None, feeds)
+        print(string_type(res, with_shape=True, with_min_max=True))
     """
 
     class IO:
@@ -172,13 +284,15 @@ class TorchOnnxEvaluator:
         opsets: Optional[Dict[str, int]] = None,
         local_functions: Optional[Dict[Tuple[str, str], "TorchOnnxEvaluator"]] = None,
         verbose: int = 0,
+        custom_kernels: Optional[Dict[Tuple[str, str], type[torch_ops.OpRunKernel]]] = None,
     ):
         self.providers = providers
         self.constants: Dict[str, torch.Tensor] = {}
-        self.kernels: List[Optional[torch_ops.OpRun]] = []
+        self.kernels: List[Optional[torch_ops.OpRunKernel]] = []
         self.functions = local_functions.copy() if local_functions else {}
         self.CPU = torch.tensor([0]).to("cpu").device
         self.verbose = verbose
+        self.custom_kernels = custom_kernels or {}
         dev = self._on_cuda(providers)
         if dev < 0:
             self.default_device = self.CPU
@@ -296,9 +410,24 @@ class TorchOnnxEvaluator:
         kernels = get_kernels()
         self.kernels.clear()
         for node in nodes:
+            kernel_kwargs = dict(verbose=max(0, self.verbose - 1))
+            opset = self.opsets[node.domain]
+            key = node.domain, node.op_type, opset
+            if key[:2] in self.custom_kernels:
+                cls = self.custom_kernels[key[:2]]
+                ags = [self.default_device] if cls.device_dependent() else []
+                kws = dict(parent=self) if cls.has_subgraphs() else {}
+                kws.update(kernel_kwargs)  # type: ignore[arg-type]
+                kernel2 = cls(node, opset, *ags, **kws)  # type: ignore[arg-type]
+                self.kernels.append(kernel2)
+                continue
+
             if (node.domain, node.op_type) in self.functions:
                 kernel = torch_ops.OpRunFunction(
-                    self.functions[node.domain, node.op_type], node, self.opsets[node.domain]
+                    self.functions[node.domain, node.op_type],
+                    node,
+                    self.opsets[node.domain],
+                    **kernel_kwargs,
                 )
                 self.kernels.append(kernel)
                 continue
@@ -308,8 +437,6 @@ class TorchOnnxEvaluator:
                 self.kernels.append(None)
                 continue
 
-            opset = self.opsets[node.domain]
-            key = node.domain, node.op_type, opset
             while key not in kernels and opset > 0:
                 opset -= 1
                 key = node.domain, node.op_type, opset
@@ -320,7 +447,8 @@ class TorchOnnxEvaluator:
             cls = kernels[key]
             ags = [self.default_device] if cls.device_dependent() else []
             kws = dict(parent=self) if cls.has_subgraphs() else {}
-            kernel2 = cls(node, opset, *ags, **kws)
+            kws.update(kernel_kwargs)  # type: ignore[arg-type]
+            kernel2 = cls(node, opset, *ags, **kws)  # type: ignore[arg-type]
             self.kernels.append(kernel2)
 
     def run(
@@ -438,7 +566,9 @@ class TorchOnnxEvaluator:
         context: Optional[Dict[str, RuntimeValue]] = None,
     ) -> Union[torch_ops.OpRunValue, Tuple[torch_ops.OpRunValue, ...]]:
         """
-        Runs the ONNX model.
+        Runs the ONNX model. The signature is different.
+        This method is called by every kernel hokding a subgraph.
+        The local variables are stored in `context`.
 
         :param args: inputs
         :param context: local context for the execution of subgraphs

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -1,4 +1,4 @@
-from ._op_run import OpRun, OpRunFunction, OpRunSequence, OpRunTensor, OpRunValue
+from ._op_run import OpRunKernel, OpRunFunction, OpRunSequence, OpRunTensor, OpRunValue
 from .access_ops import Gather_1, ScatterND_16, Slice_13
 from .binary_ops import (
     And_1,

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Optional, Union, Tuple
+from typing import Any, Dict, List, Optional, Union, Tuple
 import onnx
 import torch
 from ...api import TensorLike
@@ -167,7 +167,7 @@ class OpRunSequence(OpRunValue):
         return string_type(self.sequence, with_shape=True)
 
 
-class OpRun:
+class OpRunKernel:
     """
     Main class. Every kernel should inherit from it.
     It does not copy the proto.
@@ -185,7 +185,13 @@ class OpRun:
         """Returns True if the kernel has subgraphs."""
         return False
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+        custom_kernels: Optional[Dict[Tuple[str, str], type]] = None,
+    ):
         assert isinstance(
             node, onnx.NodeProto
         ), f"node must be a NodeProto but node is {type(node)}"
@@ -193,6 +199,8 @@ class OpRun:
         self.domain = node.domain
         self.input = node.input
         self.output = node.output
+        self.verbose = verbose
+        self.custom_kernels = custom_kernels
         if version is None:
             name = self.__class__.__name__.split("_")
             assert (
@@ -305,7 +313,7 @@ class OpRun:
         return tuple(t.to(device) for t in tensors)
 
 
-class OpRunFunction(OpRun):
+class OpRunFunction(OpRunKernel):
     """
     Defines a kernel based on a local functions.
     """
@@ -315,8 +323,9 @@ class OpRunFunction(OpRun):
         runtime: "onnx_diagnostic.reference.TorchOnnxEvaluator",  # noqa: F821
         node: onnx.NodeProto,
         version: Optional[int] = None,
+        verbose: int = 0,
     ):
-        super().__init__(node, version)
+        super().__init__(node, version, verbose=verbose)
         self.runtime = runtime
         self.input_names = runtime.input_names
 

onnx_diagnostic/reference/torch_ops/access_ops.py

@@ -1,14 +1,19 @@
 from typing import Optional
 import onnx
 import torch
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class Gather_1(OpRun):
+class Gather_1(OpRunKernel):
     "Gather"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
         axis = self.get_attribute_int(node, "axis", 0)
         assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
         self.axis = axis
@@ -21,11 +26,16 @@ class Gather_1(OpRun):
         return OpRunTensor(x.tensor[tuple(ind)])
 
 
-class ScatterND_16(OpRun):
+class ScatterND_16(OpRunKernel):
     "ScatterND"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
         self.reduction = self.get_attribute_string(node, "reduction", "none")
 
     def run(
@@ -54,7 +64,7 @@ class ScatterND_16(OpRun):
         return OpRunTensor(output)
 
 
-class Slice_13(OpRun):
+class Slice_13(OpRunKernel):
     "Slice"
 
     def run(

onnx_diagnostic/reference/torch_ops/binary_ops.py

@@ -1,8 +1,8 @@
 import torch
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class OpRunBinary(OpRun):
+class OpRunBinary(OpRunKernel):
     "Binary Op"
 
     def run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:

onnx_diagnostic/reference/torch_ops/controlflow_ops.py

@@ -1,10 +1,10 @@
 from typing import Any, Dict, Optional
 import onnx
 import torch
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class OpRunControlFlow(OpRun):
+class OpRunControlFlow(OpRunKernel):
     """Common ancestor for control flows."""
 
     @classmethod
@@ -17,8 +17,9 @@ class OpRunControlFlow(OpRun):
         node: onnx.NodeProto,
         version: Optional[int] = None,
         parent: Optional["onnx_diagnostic.reference.TorchOnnxEvaluator"] = None,  # noqa: F821
+        verbose: int = 0,
     ):
-        super().__init__(node, version)
+        super().__init__(node, version, verbose=verbose)
         assert (
             parent is not None
         ), f"parent must be specified for operator {self.__class__.__name__!r}"
@@ -30,6 +31,7 @@ class OpRunControlFlow(OpRun):
                     opsets=parent.opsets,
                     local_functions=parent.functions,
                     verbose=parent.verbose,
+                    custom_kernels=parent.custom_kernels,
                 )
                 setattr(self, att.name, rt)
 
@@ -50,8 +52,9 @@ class Loop_16(OpRunControlFlow):
         node: onnx.NodeProto,
         version: Optional[int] = None,
         parent: Optional["onnx_diagnostic.reference.TorchOnnxEvaluator"] = None,  # noqa: F821
+        verbose: int = 0,
     ):
-        super().__init__(node, version, parent)
+        super().__init__(node, version, parent, verbose=verbose)
         self.output_index = {n: i for i, n in enumerate(self.body.output_names)}
         self.N = len(self.body.input_names) - 2
         self.K = len(self.body.output_names) - self.N - 1

onnx_diagnostic/reference/torch_ops/generator_ops.py

@@ -1,10 +1,10 @@
 from typing import Optional
 import onnx
 import torch
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class Range_11(OpRun):
+class Range_11(OpRunKernel):
     """Range"""
 
     @classmethod
@@ -19,8 +19,9 @@ class Range_11(OpRun):
         node: onnx.NodeProto,
         version: Optional[int] = None,
         device: Optional[torch.device] = None,
+        verbose: int = 0,
     ):
-        super().__init__(node, version)
+        super().__init__(node, version, verbose=verbose)
         self.device = device
 
     def run(self, starts: OpRunTensor, limit: OpRunTensor, delta: OpRunTensor) -> OpRunTensor:

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -2,14 +2,19 @@ from typing import Optional, Tuple
 import onnx
 import torch
 from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class AveragePool_11(OpRun):
+class AveragePool_11(OpRunKernel):
     "AveragePool"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
         self.auto_pad = self.get_attribute_string(node, "auto_pad", "NOTSET")
         self.ceil_mode = bool(self.get_attribute_int(node, "ceil_mode", 0))
         self.count_include_pad = bool(self.get_attribute_int(node, "count_include_pad", 0))
@@ -44,11 +49,16 @@ class AveragePool_11(OpRun):
         )
 
 
-class Conv_11(OpRun):
+class Conv_11(OpRunKernel):
     "Conv"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
         self.auto_pad = self.get_attribute_string(node, "auto_pad", "NOTSET")
         self.dilations = self.get_attribute_ints(node, "dilations", None)
         self.group = self.get_attribute_int(node, "group", 1)
@@ -108,11 +118,16 @@ class Conv_11(OpRun):
         )
 
 
-class LayerNormalization_17(OpRun):
+class LayerNormalization_17(OpRunKernel):
     "LayerNormalization"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
         self.axis = self.get_attribute_int(node, "axis", -1)
         self.epsilon = self.get_attribute_float(node, "epsilon", 1e-5)
         self.stash_type = onnx_dtype_to_torch_dtype(
@@ -152,11 +167,16 @@ class LayerNormalization_17(OpRun):
         )
 
 
-class Softmax_13(OpRun):
+class Softmax_13(OpRunKernel):
     "Softmax"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
         self.axis = self.get_attribute_int(node, "axis", -1)
         assert isinstance(self.axis, int), f"Unexpected value for attribute axis={self.axis!r}"
         # this is out of spec
@@ -169,7 +189,7 @@ class Softmax_13(OpRun):
         )
 
 
-class Tanh_6(OpRun):
+class Tanh_6(OpRunKernel):
     "Tanh"
 
     def run(self, data: OpRunTensor) -> OpRunTensor:

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -2,14 +2,14 @@ from typing import Optional
 import onnx
 import torch
 from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class Cast_6(OpRun):
+class Cast_6(OpRunKernel):
     "Cast"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         to = self.get_attribute_int(node, "to", 0)
         assert isinstance(to, int), f"Unexpected value for attribute to={to!r}"
         self.to = onnx_dtype_to_torch_dtype(to)
@@ -20,11 +20,11 @@ class Cast_6(OpRun):
         return OpRunTensor(data.tensor.to(self.to))
 
 
-class CastLike_15(OpRun):
+class CastLike_15(OpRunKernel):
     "Cast"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         self.saturate = self.get_attribute_int(node, "saturate", 1)
         assert self.saturate == 1, f"saturate={self.saturate} not implemented for CastLike"
 
@@ -32,11 +32,11 @@ class CastLike_15(OpRun):
         return OpRunTensor(data.tensor.to(like.tensor.dtype))
 
 
-class Concat_1(OpRun):
+class Concat_1(OpRunKernel):
     "Concat"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         axis = self.get_attribute_int(node, "axis", 0)
         assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
         self.axis = axis
@@ -59,36 +59,36 @@ class Concat_1(OpRun):
         return OpRunTensor(torch.cat([t.tensor.to(device) for t in data], axis=self.axis))
 
 
-class NonZero_13(OpRun):
+class NonZero_13(OpRunKernel):
     "NonZero"
 
     def run(self, x: OpRunTensor) -> OpRunTensor:
         return OpRunTensor(torch.nonzero(x.tensor).T)
 
 
-class Tile_6(OpRun):
+class Tile_6(OpRunKernel):
     "Tile"
 
     def run(self, x: OpRunTensor, repeat: OpRunTensor) -> OpRunTensor:
         return OpRunTensor(torch.tile(x.tensor, repeat.as_tuple_int))
 
 
-class Transpose_1(OpRun):
+class Transpose_1(OpRunKernel):
     "Transpose"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         self.perm = self.get_attribute_ints(node, "perm", None)
 
     def run(self, data: OpRunTensor) -> OpRunTensor:
         return OpRunTensor(torch.permute(data.tensor, self.perm))
 
 
-class Trilu_14(OpRun):
+class Trilu_14(OpRunKernel):
     "Trilu"
 
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         self.upper = self.get_attribute_int(node, "upper", 1)
 
     def run(self, data: OpRunTensor, k: Optional[OpRunTensor] = None) -> OpRunTensor:
@@ -98,7 +98,7 @@ class Trilu_14(OpRun):
         return OpRunTensor(torch.tril(data.tensor, diagonal=diagonal))
 
 
-class Where_9(OpRun):
+class Where_9(OpRunKernel):
     "Where"
 
     def run(self, cond: OpRunTensor, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:

onnx_diagnostic/reference/torch_ops/reduce_ops.py

@@ -2,12 +2,12 @@ from typing import Optional, Tuple
 import onnx
 import torch
 from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
-from . import OpRun, OpRunTensor
+from . import OpRunKernel, OpRunTensor
 
 
-class ReduceOp(OpRun):
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+class ReduceOp(OpRunKernel):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         self.keepdims = bool(self.get_attribute_int(node, "keepdims", 1))
         self.noop_with_empty_axes = bool(
             self.get_attribute_int(node, "noop_with_empty_axes", 0)
@@ -28,8 +28,8 @@ class ReduceOp(OpRun):
 
 
 class ReduceOpAxes(ReduceOp):
-    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
-        super().__init__(node, version)
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
         self.axes: Tuple[int, ...] = self.get_attribute_ints(node, "axes") or tuple()