reduce

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -472,7 +472,7 @@ def test_op_reduce_min_17(self):
     def test_op_reduce_min_17_no_axes(self):
         model = oh.make_model(
             oh.make_graph(
-                [oh.make_node("ReduceMin", ["X"], ["Z"])],
+                [oh.make_node("ReduceMin", ["X"], ["Z"], keepdims=0)],
                 "dummy",
                 [oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"])],
                 [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -122,14 +122,24 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
 
     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 self.stash_type == torch.float32 and x.tensor.dtype != torch.float64:
+            xt = x.tensor
+            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,
+            )
+        else:
+            xt = x.tensor.to(self.stash_type)
+            res = torch.nn.functional.layer_norm(
+                xt,
+                xt.shape[self.axis :],
+                weight=scale.tensor.to(self.stash_type),
+                bias=None if bias is None else bias.tensor.to(self.stash_type),
+                eps=self.epsilon,
+            )
         if not self.compute_std:
             return OpRunTensor(res.to(original_dtype))
         axes = tuple(range(len(xt.shape)))[self.axis :]