fix reshape

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -1,5 +1,4 @@
 import unittest
-from typing import Optional
 import numpy as np
 import onnx
 import onnx.helper as oh
@@ -23,7 +22,7 @@ def test_kernels(self):
         kernel = ker[key]
         self.assertEqual("Add_1", kernel.__name__)
 
-    def _finalize_test(self, model, *args, atol: Optional[float] = None):
+    def _finalize_test(self, model, *args, atol: float = 0):
         onnx.checker.check_model(model)
         feeds = dict(zip([i.name for i in model.graph.input], args))
 
@@ -178,7 +177,29 @@ def test_op_reshape(self):
             opset_imports=[oh.make_opsetid("", 18)],
         )
         self._finalize_test(
-            model, torch.rand((4, 5, 6, 7), dtype=torch.float32), torch.tensor([7, 4, 6, 5])
+            model,
+            torch.rand((4, 5, 6, 7), dtype=torch.float32),
+            torch.tensor([7, 4, 6, 5], dtype=torch.int64),
+        )
+
+    def test_op_reshape_zero(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Reshape", ["X", "shape"], ["Y"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c", "d"]),
+                    oh.make_tensor_value_info("shape", TINT64, ["f"]),
+                ],
+                [oh.make_tensor_value_info("Y", TFLOAT, ["d", "a", "c", "b"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand((4, 5, 6, 7), dtype=torch.float32),
+            torch.tensor([7, 4, 0, 5], dtype=torch.int64),
         )
 
     def test_op_matmul(self):

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -3,4 +3,4 @@
 from .binary_ops import Add_1, Div_1, MatMul_1, Mul_1, Sub_1
 from .other_ops import Cast_6, Concat_1, Transpose_1
 from .nn_ops import Softmax_13
-from .shape_ops import Reshape_5, Shape_15, Squeeze_13, Unsqueeze_13
+from .shape_ops import Reshape_14, Shape_15, Squeeze_13, Unsqueeze_13

onnx_diagnostic/reference/torch_ops/shape_ops.py

@@ -4,10 +4,21 @@
 from . import OpRun, OpRunValue
 
 
-class Reshape_5(OpRun):
+class Reshape_14(OpRun):
     "Reshape"
 
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.allowzero = self.get_attribute_int(node, "allowzero", 0)
+
     def run(self, data: OpRunValue, shape: OpRunValue) -> OpRunValue:
+        ishape = shape.as_tuple_int
+        if not self.allowzero and 0 in ishape:
+            xshape = data.tensor.shape
+            new_shape = []
+            for i, s in enumerate(ishape):
+                new_shape.append(xshape[i] if s == 0 else s)
+            return OpRunValue(data.tensor.reshape(new_shape))
         return OpRunValue(data.tensor.reshape(shape.as_tuple_int))