enhance cond_rewriter by allowing different output shapes

Author: lucienwang1009

tests/test_cond.py

@@ -12,7 +12,7 @@
 import tensorflow as tf
 
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main
+from common import unittest_main, check_opset_min_version
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -267,6 +267,51 @@ def case_graph():
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port)
 
+    @check_opset_min_version(9, "")
+    def test_cond_with_different_output_shape(self):
+        input_shape = (10, 5, 20)
+        inputs = tf.placeholder(tf.float32, input_shape, name="input")
+
+        shape = tf.placeholder(tf.int32, (len(input_shape),), name="shape")
+        # cheat onnx shape inference
+        inputs = tf.reshape(inputs, shape)
+
+        def pad_tensor(t, length):
+          """Pads the input tensor with 0s along the first dimension up to the length.
+
+          Args:
+            t: the input tensor, assuming the rank is at least 1.
+            length: a tensor of shape [1]  or an integer, indicating the first dimension
+              of the input tensor t after padding, assuming length <= t.shape[0].
+
+          Returns:
+            padded_t: the padded tensor, whose first dimension is length. If the length
+              is an integer, the first dimension of padded_t is set to length
+              statically.
+          """
+          t_rank = tf.rank(t)
+          t_shape = tf.shape(t)
+          t_d0 = t_shape[0]
+          pad_d0 = tf.expand_dims(length - t_d0, 0)
+          pad_shape = tf.cond(
+              # shape is [3], depending on input shape
+              tf.greater(t_rank, 1), lambda: tf.concat([pad_d0, t_shape[1:]], 0),
+              # shape is always [1]
+              lambda: tf.expand_dims(length - t_d0, 0))
+          padded_t = tf.concat([t, tf.zeros(pad_shape, dtype=t.dtype)], 0)
+          return padded_t
+
+        output = pad_tensor(inputs, 20)
+        _ = tf.identity(output, name="output")
+        input_names_with_port = ["input:0", "shape:0"]
+        feed_dict = {
+            "input:0": np.ones(input_shape, dtype=np.float32),
+            "shape:0": np.array(input_shape, dtype=np.int32)
+        }
+
+        output_names_with_port = ["output:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/rewriter/cond_rewriter.py

@@ -106,16 +106,21 @@ def _get_output_shape_dtype(self, cond_context):
             true_output = cond_context.true_branch_context.output[i]
             false_output = cond_context.false_branch_context.output[i]
             true_shape = self.g.get_shape(true_output)
+            utils.make_sure(true_shape is not None, "Shape of {} is None".format(true_output))
+            true_rank  = len(true_shape)
             true_dtype = self.g.get_dtype(true_output)
             false_shape = self.g.get_shape(false_output)
+            utils.make_sure(false_shape is not None, "Shape of {} is None".format(false_output))
+            false_rank = len(false_shape)
             false_dtype = self.g.get_dtype(false_output)
-            if not utils.are_shapes_compatible(true_shape, false_shape):
+            # just require rank is equal
+            if true_rank != false_rank:
                 raise RuntimeError(
-                    "the shape of outputs {} and {} mismatch: {}, {}".format(
+                    "the rank of outputs {} and {} mismatch: {}, {}".format(
                         true_output,
                         false_output,
-                        true_shape,
-                        false_shape
+                        true_rank,
+                        false_rank
                     )
                 )
             if true_dtype != false_dtype:
@@ -127,10 +132,7 @@ def _get_output_shape_dtype(self, cond_context):
                         false_dtype
                     )
                 )
-            # in tf, the shape of different branched can be different,
-            # for example output shape of branch A can be [-1] while branch B can be [1].
-            # Under this case, we should set output shape to be [-1]
-            output_shapes.append(utils.create_vague_shape_like(utils.merge_shapes(true_shape, false_shape)))
+            output_shapes.append(utils.create_vague_shape_like(true_shape))
             output_dtypes.append(true_dtype)
         return output_shapes, output_dtypes