f

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -22,13 +22,20 @@ def test_kernels(self):
         kernel = ker[key]
         self.assertEqual("Add_1", kernel.__name__)
 
-    def _finalize_test(self, model, *args, atol: float = 0):
+    def _finalize_test(self, model, *args, atol: float = 0, use_ort: bool = False):
         onnx.checker.check_model(model)
         feeds = dict(zip([i.name for i in model.graph.input], args))
+        feeds_numpy = {k: v.numpy() for k, v in feeds.items()}
 
-        expected = ExtendedReferenceEvaluator(model).run(
-            None, {k: v.numpy() for k, v in feeds.items()}
-        )
+        if use_ort:
+            import onnxruntime
+
+            sess = onnxruntime.InferenceSession(
+                model.SerializeToString(), providers=["CPUExecutionProvider"]
+            )
+            expected = sess.run(None, feeds_numpy)
+        else:
+            expected = ExtendedReferenceEvaluator(model).run(None, feeds_numpy)
         rt = TorchEvaluator(model)
         got = rt.run(None, feeds)
         self.assertEqualAny(expected, [g.detach().numpy() for g in got], atol=atol)
@@ -453,6 +460,52 @@ def test_op_reduce_sum(self):
             atol=1e-5,
         )
 
+    def test_op_where(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Greater", ["X", "Y"], ["cond"]),
+                    oh.make_node("Where", ["cond", "X", "Y"], ["Z"]),
+                ],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            torch.rand(3, 4, 5, dtype=torch.float32),
+        )
+
+    def test_op_layer_normalization(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("LayerNormalization", ["X", "W", "B"], ["Z"], axis=-1)],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"]),
+                    oh.make_tensor_value_info("W", TFLOAT, []),
+                    oh.make_tensor_value_info("B", TFLOAT, []),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            torch.abs(torch.rand(5, dtype=torch.float32)),
+            torch.rand(5, dtype=torch.float32),
+            use_ort=True,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -13,8 +13,8 @@
     Or_1,
     Sub_1,
 )
-from .other_ops import Cast_6, Concat_1, Transpose_1
-from .nn_ops import Softmax_13, Tanh_6
+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/_op_run.py

@@ -80,6 +80,20 @@ def _find_attribute(self, node: onnx.NodeProto, name: str):
                 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]:

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -5,6 +5,36 @@
 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)
+        )
+        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 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 res.to(original_dtype), mean, x_inv_std_dev
+
+
 class Softmax_13(OpRun):
     "Softmax"
 

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -40,3 +40,10 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = 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))