Add Shape op and other ops (#118)

* Add Shape op

* shape

* op

* add concat

* add gather

* softmax

* fix reshape

* tanh

* add equal

* unary

* reducemin

* reduce

* ind

* disable

* f

* fix layer norm

* mypy

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -69,7 +69,7 @@ def test_validate_model_export(self):
         self.assertIsInstance(summary, dict)
         self.assertIsInstance(data, dict)
 
-    @requires_torch("2.7")
+    @requires_torch("2.8.99")
     @hide_stdout()
     @ignore_warnings(FutureWarning)
     def test_validate_model_onnx_dynamo_ir(self):

_unittests/ut_torch_models/test_test_helpers.py

@@ -907,11 +907,13 @@ def assertEqualArray(
         except AssertionError as e:
             expected_max = numpy.abs(expected).max()
             expected_value = numpy.abs(value).max()
+            te = expected.astype(int) if expected.dtype == numpy.bool_ else expected
+            tv = value.astype(int) if value.dtype == numpy.bool_ else value
             rows = [
                 f"{msg}\n{e}" if msg else str(e),
                 f"expected max value={expected_max}",
                 f"expected computed value={expected_value}\n",
-                f"ratio={expected / value}\ndiff={expected - value}",
+                f"ratio={te / tv}\ndiff={te - tv}",
             ]
             raise AssertionError("\n".join(rows))  # noqa: B904
 

onnx_diagnostic/ext_test_case.py

@@ -172,16 +172,23 @@ def run(
                 inputs = [(self.runtime_info[i].value if i else None) for i in kernel.input]
                 res = kernel.run(*inputs)
                 if isinstance(res, tuple):
+                    # outputs
+                    assert all(isinstance(o, torch_ops.OpRunValue) for o in res), (
+                        f"Unexpected output type {[type(o) for o in res]} "
+                        f"for kernel {type(kernel)}."
+                    )
                     for name, t in zip(kernel.output, res):
                         self.runtime_info[name].set_value(t)
                 else:
+                    assert isinstance(
+                        res, torch_ops.OpRunValue
+                    ), f"Unexpected output type {type(res)} for kernel {type(kernel)}."
                     self.runtime_info[kernel.output[0]].set_value(res)
 
             # free intermediate results
             for name in self.last_used[it]:
                 self.runtime_info[name].clean_value()
 
-        # outputs
         res = [self.runtime_info[o].value.tensor for o in outputs]  # type: ignore[assignment, union-attr]
 
         # clean previous execution

onnx_diagnostic/reference/torch_evaluator.py

@@ -1,4 +1,20 @@
 from ._op_run import OpRun, OpRunValue
-from .access_ops import Slice_13
-from .binary_ops import Add_1, Div_1, Mul_1, Sub_1
-from .shape_ops import Squeeze_13
+from .access_ops import Gather_1, Slice_13
+from .binary_ops import (
+    And_1,
+    Add_1,
+    Div_1,
+    Greater_1,
+    GreaterOrEqual_1,
+    Less_1,
+    LessOrEqual_1,
+    MatMul_1,
+    Mul_1,
+    Or_1,
+    Sub_1,
+)
+from .nn_ops import LayerNormalization_17, Softmax_13, Tanh_6
+from .other_ops import Cast_6, Concat_1, Transpose_1, Where_9
+from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_18
+from .shape_ops import Reshape_14, Shape_15, Squeeze_13, Unsqueeze_13
+from .unary_ops import Neg_1, Not_1, Reciprocal_1

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -73,3 +73,51 @@ def run(self, *args) -> Union[OpRunValue, Tuple[OpRunValue, ...]]:
         raise NotImplementedError(
             f"Method run is not implemented for kernel {self.__class__.__name__!r}"
         )
+
+    def _find_attribute(self, node: onnx.NodeProto, name: str):
+        for att in node.attribute:
+            if att.name == name:
+                return att
+        return None
+
+    def get_attribute_float(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[float] = None
+    ) -> Optional[float]:
+        """
+        Returns an attribute as an int.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else float(att.f)
+
+    def get_attribute_int(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[int] = None
+    ) -> Optional[int]:
+        """
+        Returns an attribute as an int.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else int(att.i)
+
+    def get_attribute_ints(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[Tuple[int, ...]] = None
+    ) -> Optional[Tuple[int, ...]]:
+        """
+        Returns an attribute as an int.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else tuple(att.ints)

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -1,7 +1,26 @@
 from typing import Optional
+import onnx
+import torch
 from . import OpRun, OpRunValue
 
 
+class Gather_1(OpRun):
+    "Gather"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        axis = self.get_attribute_int(node, "axis", 0)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+
+    def run(self, x, indices):
+        if indices.tensor.numel() == 0:
+            return torch.empty((0,), dtype=x.tensor.dtype, device=x.tensor.device)
+        ind = [slice(0, s) for s in x.shape]
+        ind[self.axis] = indices.tensor
+        return OpRunValue(x.tensor[tuple(ind)])
+
+
 class Slice_13(OpRun):
     "Slice"
 

onnx_diagnostic/reference/torch_ops/access_ops.py

@@ -1,25 +1,81 @@
 from . import OpRun, OpRunValue
 
 
+class And_1(OpRun):
+    """And"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor & y.tensor)
+
+
 class Add_1(OpRun):
     """Add"""
 
     def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
         return OpRunValue(x.tensor + y.tensor)
 
 
+class Div_1(OpRun):
+    """Div"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor / y.tensor)
+
+
+class Equal_1(OpRun):
+    """Equal"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor == y.tensor)
+
+
+class Greater_1(OpRun):
+    """Greater"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor > y.tensor)
+
+
+class GreaterOrEqual_1(OpRun):
+    """GreaterOrEqual"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor >= y.tensor)
+
+
+class Less_1(OpRun):
+    """Less"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor < y.tensor)
+
+
+class LessOrEqual_1(OpRun):
+    """LessOrEqual"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor <= y.tensor)
+
+
+class MatMul_1(OpRun):
+    """MatMul"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor @ y.tensor)
+
+
 class Mul_1(OpRun):
     """Mul"""
 
     def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
         return OpRunValue(x.tensor * y.tensor)
 
 
-class Div_1(OpRun):
-    """Div"""
+class Or_1(OpRun):
+    """Or"""
 
     def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
-        return OpRunValue(x.tensor / y.tensor)
+        return OpRunValue(x.tensor | y.tensor)
 
 
 class Sub_1(OpRun):

onnx_diagnostic/reference/torch_ops/binary_ops.py

@@ -0,0 +1,59 @@
+from typing import Optional
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRun, OpRunValue
+
+
+class LayerNormalization_17(OpRun):
+    "LayerNormalization"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        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)  # type: ignore[arg-type]
+        )
+        self.compute_std = len(node.output) > 1
+
+    def run(self, x, scale, bias=None):
+        original_dtype = x.dtype
+        xt = x.tensor.to(self.stash_type)
+        res = torch.nn.functional.layer_norm(
+            xt,
+            xt.shape[self.axis :],
+            weight=scale.tensor,
+            bias=None if bias is None else bias.tensor,
+            eps=self.epsilon,
+        )
+        if not self.compute_std:
+            return OpRunValue(res.to(original_dtype))
+        axes = tuple(range(len(xt.shape)))[self.axis :]
+        mean, var = torch.var(xt, dim=axes, keepdim=False)
+        x_inv_std_dev = torch.reciprocal(torch.sqrt(var + self.epsilon))
+        return OpRunValue(res.to(original_dtype)), OpRunValue(mean), OpRunValue(x_inv_std_dev)
+
+
+class Softmax_13(OpRun):
+    "Softmax"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        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
+        stash_type = self.get_attribute_int(node, "stash_type", None)
+        self.stash_type = None if stash_type is None else onnx_dtype_to_torch_dtype(stash_type)
+
+    def run(self, data: OpRunValue) -> OpRunValue:
+        return OpRunValue(
+            torch.nn.functional.softmax(data.tensor, dim=self.axis, dtype=self.stash_type)
+        )
+
+
+class Tanh_6(OpRun):
+    "Tanh"
+
+    def run(self, data: OpRunValue) -> OpRunValue:
+        return OpRunValue(torch.nn.functional.tanh(data.tensor))

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -0,0 +1,49 @@
+from typing import Optional
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRun, OpRunValue
+
+
+class Cast_6(OpRun):
+    "Cast"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        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)
+
+    def run(self, data: OpRunValue) -> OpRunValue:
+        return OpRunValue(data.tensor.to(self.to))
+
+
+class Concat_1(OpRun):
+    "Concat"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        axis = self.get_attribute_int(node, "axis", 0)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+
+    def run(self, *data: OpRunValue) -> OpRunValue:
+        return OpRunValue(torch.cat([t.tensor for t in data], axis=self.axis))
+
+
+class Transpose_1(OpRun):
+    "Transpose"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.perm = self.get_attribute_ints(node, "perm", None)
+
+    def run(self, data: OpRunValue) -> OpRunValue:
+        return OpRunValue(torch.permute(data.tensor, self.perm))
+
+
+class Where_9(OpRun):
+    "Where"
+
+    def run(self, cond: OpRunValue, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(torch.where(cond.tensor, x.tensor, y.tensor))