Remove double reshape (#1495)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1341,6 +1341,17 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, check_shape=True)
 
+    def test_reshape_reshape(self):
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2))
+        def func(x):
+            shape = tf.constant([1, 4])
+            shape_2 = tf.constant([4, 1])
+            x_ = tf.reshape(x, shape)
+            x_ = tf.reshape(x_, shape_2)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val},
+                            graph_validator=lambda g: check_op_count(g, "Reshape", 1, disabled=False))
+
     @check_opset_min_version(6, "cast")
     def test_reshape_int(self):
         x_val = np.array([1, 2, 3, 4], dtype=np.int32).reshape((2, 2))

tf2onnx/optimizer/back_to_back_optimizer.py

@@ -17,6 +17,8 @@
 
 
 def _register_func(op_type):
+    if not isinstance(op_type, tuple):
+        op_type = (op_type,)
     def _internal_fun(func):
         _func_map[op_type] = func
         return func
@@ -252,3 +254,18 @@ def _optimize_conv_batchnorm_fusion(g, node, consumer_nodes):
         g.set_shape(node2_output[0], node2_shape)
         g.set_dtype(node2_output[0], node2_dtype)
         return []
+
+    @staticmethod
+    @_register_func('Reshape')
+    def _optimize_reshape_reshape(g, node, consumer_nodes):
+        """remove sequential reshape nodes"""
+        if node.type != 'Reshape' or len(consumer_nodes) != 1:
+            return []
+
+        node2 = consumer_nodes[0]
+        if node2.type != 'Reshape':
+            return []
+
+        g.replace_inputs(node2, [node.input[0], node2.input[1]])
+        g.remove_node(node.name)
+        return []