code refactor according to reviewer's comments

Author: zhijxu-MS

tests/test_backend.py

@@ -201,33 +201,6 @@ def test_atrig_ops(self):
             _ = tf.identity(op_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    def test_map_fn(self):
-        def fn0(elem):
-            res = elem + elem * elem
-            return res
-
-        def fn1(elem):
-            res = elem[0] * elem[1] + elem[0]
-            return res
-
-        x_val = 100 * np.random.random_sample([2, 10]).astype(np.float32)
-        y_val = 100 * np.random.random_sample([2, 10]).astype(np.float32)
-
-        # test fn0
-        x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
-        res_ = tf.map_fn(fn0, x, dtype=tf.float32)
-        _ = tf.identity(res_, name=_TFOUTPUT1)
-        self._run_test_case([_OUTPUT1], {_INPUT: x_val}, rtol=0)
-        tf.reset_default_graph()
-
-        # test fn1
-        x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
-        y = tf.placeholder(tf.float32, shape=y_val.shape, name=_TFINPUT1)
-        res_ = tf.map_fn(fn1, (x, y), dtype=tf.float32)
-        _ = tf.identity(res_, name=_TFOUTPUT1)
-        self._run_test_case([_OUTPUT1], {_INPUT: x_val, _INPUT1: y_val}, rtol=0)
-        tf.reset_default_graph()
-
     @unittest.skipIf(BACKEND in ["caffe2"], "not supported correctly in caffe2")
     @unittest.skipIf(*support_op_conversion_since(7, "multinomial"))
     def test_multinomial(self):

tests/test_loops.py

@@ -16,6 +16,7 @@
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
 
+
 class LoopTests(Tf2OnnxBackendTestBase):
 
     def test_simple_while_loop(self):
@@ -31,7 +32,6 @@ def test_simple_while_loop(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
-
     def test_simple_while_loop_2(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         c = lambda i: tf.logical_and(tf.less(i, 10), tf.greater_equal(i, 3))
@@ -45,7 +45,6 @@ def test_simple_while_loop_2(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
-
     def test_while_loop_with_ta_write(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         output_ta = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
@@ -68,7 +67,6 @@ def b(i, out_ta):
         output_names_with_port = ["output:0", "i:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
-
     def test_while_loop_with_ta_read(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         inputs = tf.placeholder(tf.float32, (10,), name="input_2")
@@ -82,6 +80,7 @@ def test_while_loop_with_ta_read(self):
         c = lambda i, *_: tf.logical_and(tf.less(i, 10), i >= 0)
         res = tf.constant(0.)
         res2 = tf.constant(1.)
+
         def b(i, res, res2):
             new_i = tf.add(i, 1)
             x = input_ta.read(i)
@@ -113,6 +112,7 @@ def test_while_loop_with_ta_read_reference_outer_input_directly(self):
         c = lambda i, *_: tf.logical_and(tf.less(i, 10), i >= 0)
         res = tf.constant(0.)
         res2 = tf.constant(1.)
+
         def b(i, res, res2):
             new_i = tf.add(i, 1)
             x = input_ta.read(i)
@@ -132,7 +132,6 @@ def b(i, res, res2):
         output_names_with_port = ["i:0", "x:0", "y:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
-
     def test_while_loop_with_ta_read_and_write(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         inputs = tf.placeholder(tf.float32, (10,), name="input_2")
@@ -160,5 +159,42 @@ def b(i, out_ta):
         output_names_with_port = ["i:0", "output_ta:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
+    def test_map_fn(self):
+        def fn0(elem):
+            res = elem + elem * elem
+            return res
+
+        def fn1(elem):
+            res = elem[0] * elem[1] + elem[0]
+            return res
+
+        x_val = 100 * np.random.random_sample([2, 10]).astype(np.float32)
+        y_val = 100 * np.random.random_sample([2, 10]).astype(np.float32)
+
+        # test fn0
+        x = tf.placeholder(tf.float32, shape=x_val.shape, name="input_0")
+        x_ = tf.identity(x)
+        res_ = tf.map_fn(fn0, x_, dtype=tf.float32)
+        _ = tf.identity(res_, name="output_0")
+        feed_dict = {"input_0:0": x_val}
+        input_names_with_port = ["input_0:0"]
+        output_names_with_port = ["output_0:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-5)
+        tf.reset_default_graph()
+
+        # test fn1
+        x = tf.placeholder(tf.float32, shape=x_val.shape, name="input_0")
+        y = tf.placeholder(tf.float32, shape=y_val.shape, name="input_1")
+        x_ = tf.identity(x)
+        y_ = tf.identity(y)
+        res_ = tf.map_fn(fn1, (x_, y_), dtype=tf.float32)
+        _ = tf.identity(res_, name="output_0")
+        feed_dict = {"input_0:0": x_val, "input_1:0": y_val}
+        input_names_with_port = ["input_0:0", "input_1:0"]
+        output_names_with_port = ["output_0:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-5)
+        tf.reset_default_graph()
+
+
 if __name__ == '__main__':
     Tf2OnnxBackendTestBase.trigger(LoopTests)

tf2onnx/shape_inference.py

@@ -10,6 +10,7 @@
 from __future__ import unicode_literals
 import logging
 from onnx import onnx_pb
+from tf2onnx import utils
 
 # pylint: disable=logging-not-lazy,missing-docstring,consider-swap-variables
 
@@ -232,17 +233,9 @@ def infer_output_shapes_with_partial_inputs(g, node):
             return False
 
         # find TensorArrayWrite
-        tensor_array_consumers = g.find_output_consumers(handle_node.output[0])
-        tensor_array_write_found = False
-        for i in tensor_array_consumers:
-            if tensor_array_write_found:
-                break
-            consumer_nodes = g.find_output_consumers(i.output[0])
-            for j in consumer_nodes:
-                if i.type == "Enter" and j.type == "TensorArrayWriteV3":
-                    tensor_array_write_node = j
-                    tensor_array_write_found = True
-                    break
+        tensor_array_write_node = _find_tensorarray_write(g, handle_node)
+        if not tensor_array_write_node:
+            return False
         # get TensorArray shape from input tensor of the found TensorArrayWrite node
         value_node = tensor_array_write_node.inputs[2]
         shape = g.get_shape(value_node.output[0])
@@ -347,3 +340,16 @@ def broadcast_shape_inference(shape_0, shape_1):
             return None
         i -= 1
     return new_shape
+
+
+def _find_tensorarray_write(graph, node):
+    utils.make_sure(node.type == "TensorArrayV3", "node should be tensorarray")
+
+    tensor_array_consumers = graph.find_output_consumers(node.output[0])
+    for i in tensor_array_consumers:
+        if i.type == "Enter":
+            consumer_nodes = graph.find_output_consumers(i.output[0])
+            for j in consumer_nodes:
+                if j.type == "TensorArrayWriteV3":
+                    return j
+    return None