resolve comments

Author: lucienwang1009

tests/test_loops.py

@@ -20,24 +20,6 @@
 
 class LoopTests(Tf2OnnxBackendTestBase):
 
-    @check_tf_min_version("1.9")
-    def test_simple_while_loop_var_shape(self):
-        # test for while_loop with variant shape variables
-        # may not meet ONNX Loop spec
-        i = tf.placeholder(tf.int32, (1), name="input_1")
-        const = tf.constant(np.array([2], dtype=np.int32))
-
-        c = lambda i: tf.reduce_all(tf.shape(i) < 10)
-        b = lambda i: tf.concat([i, const], 0)
-        r = tf.while_loop(c, b, [i], shape_invariants=[tf.TensorShape([None])])
-
-        _ = tf.identity(r, name="output")
-        input_names_with_port = ["input_1:0"]
-        feed_dict = {"input_1:0": np.array([0], 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)
-
     def test_simple_while_loop(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         c = lambda i: tf.less(i, 10)
@@ -214,6 +196,24 @@ def fn1(elem):
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-5)
         tf.reset_default_graph()
 
+    @check_tf_min_version("1.9")
+    def test_simple_while_loop_var_shape(self):
+        # test for while_loop with variant shape variables
+        # may not meet ONNX Loop spec
+        i = tf.placeholder(tf.int32, (1), name="input_1")
+        const = tf.constant(np.array([2], dtype=np.int32))
+
+        c = lambda i: tf.reduce_all(tf.shape(i) < 10)
+        b = lambda i: tf.concat([i, const], 0)
+        r = tf.while_loop(c, b, [i], shape_invariants=[tf.TensorShape([None])])
+
+        _ = tf.identity(r, name="output")
+        input_names_with_port = ["input_1:0"]
+        feed_dict = {"input_1:0": np.array([0], 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()

tests/test_tf_shape_inference.py

@@ -68,9 +68,9 @@ def _run_test_case(self, input_names_with_port, output_names_with_port):
     def _compare_shape_for_op(self, op1, op2):
         """Align outputs of op2 to op1."""
         for out1, out2 in zip(op1.outputs, op2.outputs):
-            expected_shape = utils.get_shape_from_tf_output(out1)
+            expected_shape = utils.get_tf_tensor_shape(out1)
             if out1 is not None:
-                actual_shape = utils.get_shape_from_tf_output(out2)
+                actual_shape = utils.get_tf_tensor_shape(out2)
                 self.assertTrue(utils.are_shapes_compatible(expected_shape, actual_shape))
 
     def test_while_loop_with_ta_read_and_write(self):

tf2onnx/rewriter/bigru_rewriter.py

@@ -13,7 +13,7 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.utils import is_reverse_op
+from tf2onnx.utils import is_tf_reverse_op
 from tf2onnx.rewriter.bilstm_rewriter import slice_bilstm_for_original_lstm_consumers,\
      get_reverse_nodes_after_y_output, get_np_val_for_const, _process_single_init_node
 
@@ -143,7 +143,7 @@ def rewrite_bidirectional_grus(g, ops):
             input_id = temp.input[0]
             temp = temp.inputs[0]
 
-        if is_reverse_op(temp):
+        if is_tf_reverse_op(temp):
             input_id = temp.input[0]
             is_backward_gru = True
 

tf2onnx/rewriter/bilstm_rewriter.py

@@ -13,7 +13,7 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.utils import is_reverse_op
+from tf2onnx.utils import is_tf_reverse_op
 from tf2onnx.graph_builder import GraphBuilder
 
 logger = logging.getLogger(__name__)
@@ -186,7 +186,7 @@ def rewrite_bidirectional_lstms(g, ops):
             input_id = temp.input[0]
             temp = temp.inputs[0]
 
-        if is_reverse_op(temp):
+        if is_tf_reverse_op(temp):
             input_id = temp.input[0]
             is_backward_lstm = True
 
@@ -220,13 +220,13 @@ def get_reverse_nodes_after_y_output(g, lstm_bw):
         if len(trans_nodes) == 1:
             if trans_nodes[0].type == "Transpose":
                 reverse_nodes = g.find_output_consumers(trans_nodes[0].output[0])
-            elif is_reverse_op(trans_nodes[0]):
+            elif is_tf_reverse_op(trans_nodes[0]):
                 reverse_nodes = trans_nodes
             else:
                 logger.debug("not found reverse op, unexpected")
                 return None
 
-            are_all_reverse = all([is_reverse_op(r_op) for r_op in reverse_nodes])
+            are_all_reverse = all([is_tf_reverse_op(r_op) for r_op in reverse_nodes])
             if are_all_reverse:
                 return reverse_nodes
 

tf2onnx/rewriter/loop_rewriter_base.py

@@ -12,8 +12,8 @@
 from collections import OrderedDict
 from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
-from tf2onnx.utils import is_loopcond_op, is_tensor_array_op
-from tf2onnx.utils import is_tensor_array_gather_op, is_tensor_array_write_op
+from tf2onnx.utils import is_tf_loopcond_op, is_tf_tensor_array_op
+from tf2onnx.utils import is_tf_tensor_array_gather_op, is_tf_tensor_array_write_op
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT
 from tf2onnx.utils import TensorValueInfo
 
@@ -195,7 +195,7 @@ def rewrite(self, context):
     def run_internal(self):
         loopcond_ops = []
         for op in self.g.get_nodes():
-            if is_loopcond_op(op):
+            if is_tf_loopcond_op(op):
                 loopcond_ops.append(op)
 
         # self.g.get_nodes may change inside this loop so that we parse all LoopCond first
@@ -316,7 +316,7 @@ def _cut_off_connection_for_cell(self, context):
 
             if val.is_tensor_array:
                 # connect NextIteration to an invalid node, to cut off an ending node of the cell.
-                ta_write_nodes = [n for n in self.g.get_nodes() if is_tensor_array_write_op(n)]
+                ta_write_nodes = [n for n in self.g.get_nodes() if is_tf_tensor_array_write_op(n)]
                 self.g.replace_all_inputs(ta_write_nodes, val.next_iteration_input.id, INVALID_INPUT_ID)
             else:
                 # connect NextIteration to an invalid node, to cut off an ending node of the cell.
@@ -376,19 +376,19 @@ def _get_loop_var_from_switch(self, switch_node):
 
         is_ta = False
         ta_index_id = None
-        if is_tensor_array_op(self.g.get_node_by_output(target_node_input_id)):
+        if is_tf_tensor_array_op(self.g.get_node_by_output(target_node_input_id)):
             is_ta = True
 
             ta_write_node = self.g.get_node_by_output(last_iteration_output_id)
-            utils.make_sure(is_tensor_array_write_op(ta_write_node), "ta nextiteration is not following ta write op")
+            utils.make_sure(is_tf_tensor_array_write_op(ta_write_node), "ta nextiteration is not following ta write op")
             last_iteration_output_id = ta_write_node.input[2]
             ta_index_id = ta_write_node.input[1]
 
             # here we parse patterns generated by
             # ta.write(), then ta.stack(), because this is the most frequent usage pattern.
             if exit_output_id:
                 exit_consumers = self.g.find_output_consumers(exit_output_id)
-                ta_gather_node = [n for n in exit_consumers if is_tensor_array_gather_op(n)][0]
+                ta_gather_node = [n for n in exit_consumers if is_tf_tensor_array_gather_op(n)][0]
 
                 # update exit output id, treat the gather output as ta's output
                 exit_output_id = ta_gather_node.output[0]

tf2onnx/rewriter/lstm_rewriter.py

@@ -14,7 +14,7 @@
 from tf2onnx import utils
 from tf2onnx.graph_builder import GraphBuilder
 from tf2onnx.rewriter.rnn_utils import RNNUnitType, RnnWeight, get_weights_from_const_node
-from tf2onnx.utils import is_concat_op, is_slice_op
+from tf2onnx.utils import is_tf_concat_op, is_tf_slice_op
 
 from tf2onnx.rewriter.unit_rnn_rewriter_base import UnitRnnRewriterBase
 
@@ -169,8 +169,8 @@ def _ct_ht_shared_variable_finder(self, context):
         lstm_cell = context.cell_match
         ct = lstm_cell.get_op("ct").output[0]
         ht = lstm_cell.get_op("ht").output[0]
-        ct_concat = [c for c in self.g.find_output_consumers(ct) if is_concat_op(c)]
-        ht_concat = [c for c in self.g.find_output_consumers(ht) if is_concat_op(c)]
+        ct_concat = [c for c in self.g.find_output_consumers(ct) if is_tf_concat_op(c)]
+        ht_concat = [c for c in self.g.find_output_consumers(ht) if is_tf_concat_op(c)]
         if len(ct_concat) != 1 or len(ht_concat) != 1 or ct_concat[0] != ht_concat[0]:
             logger.debug("failed to find ct-ht concat")
             return None
@@ -179,8 +179,8 @@ def _ct_ht_shared_variable_finder(self, context):
         consumers = []
         ct_identity_consumer = lstm_cell.get_op("ct_identity_consumer")
         ht_identity_consumer = lstm_cell.get_op("xh")
-        ct_slice = [c for c in ct_identity_consumer.inputs if is_slice_op(c)]
-        ht_slice = [c for c in ht_identity_consumer.inputs if is_slice_op(c)]
+        ct_slice = [c for c in ct_identity_consumer.inputs if is_tf_slice_op(c)]
+        ht_slice = [c for c in ht_identity_consumer.inputs if is_tf_slice_op(c)]
         if len(ct_slice) != 1 or len(ht_slice) != 1:
             logger.debug("failed to find slice op before identity consumers")
             return None

tf2onnx/rewriter/unit_rnn_rewriter_base.py

@@ -14,7 +14,7 @@
 from tf2onnx.rewriter.loop_rewriter_base import LoopRewriterBase, Context
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT, get_pattern, \
     get_rnn_scope_name, parse_rnn_loop, seq_len_pattern
-from tf2onnx.utils import is_select_op, is_tensor_array_write_op
+from tf2onnx.utils import is_tf_select_op, is_tf_tensor_array_write_op
 from tf2onnx.graph_matcher import GraphMatcher
 
 
@@ -210,7 +210,7 @@ def find_sequence_length_node(self, context):
         # get any state variable
         state_variable = list(context.state_variables.values())[0]
         next_iter_input_node = self.g.get_node_by_output(state_variable.next_iteration_input.id)
-        if not is_select_op(next_iter_input_node):
+        if not is_tf_select_op(next_iter_input_node):
             logger.debug("no sequence length node is given")
             return None
         matcher = GraphMatcher(seq_len_pattern)
@@ -302,10 +302,10 @@ def _find_state_variable_with_select(self, context,
         # find all select not followed by TensorArrayWrite
         select = []
         for c in self.g.find_output_consumers(next_iteration_input):
-            if not is_select_op(c):
+            if not is_tf_select_op(c):
                 continue
             out_ta_writer = [
-                o for o in self.g.find_output_consumers(c.output[0]) if is_tensor_array_write_op(o)
+                o for o in self.g.find_output_consumers(c.output[0]) if is_tf_tensor_array_write_op(o)
             ]
             if out_ta_writer:
                 continue