Fix pylint errors and unit tests

- fixed pylint errors
- fixed some broken for-loop logic

Author: buddhapuneeth

tests/test_stacked_lstm.py

@@ -5,10 +5,14 @@
 
 from tensorflow.python.ops import init_ops
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main, check_lstm_count, skip_tf2
+from common import unittest_main, check_lstm_count
 
 from tf2onnx.tf_loader import is_tf2
 
+
+# pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test,cell-var-from-loop
+# pylint: disable=invalid-name
+
 if is_tf2():
     LSTMCell = tf.compat.v1.nn.rnn_cell.LSTMCell
     MultiRNNCell = tf.compat.v1.nn.rnn_cell.MultiRNNCell
@@ -19,6 +23,7 @@
     MultiRNNCell = tf.contrib.rnn.MultiRNNCell
     dynamic_rnn = tf.nn.dynamic_rnn
 
+
 class LSTMLayeredTests(Tf2OnnxBackendTestBase):
     def test_layered_lstm(self):
         units = 5
@@ -52,5 +57,6 @@ def lstm_cell():
         self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06,
                            graph_validator=lambda g: check_lstm_count(g, 2))
 
+
 if __name__ == '__main__':
-    unittest_main()
+    unittest_main()

tf2onnx/rewriter/lstm_rewriter.py

@@ -40,14 +40,14 @@ def find_cell(self, context):
             cell_match = self._match_cell(context, cell_type)
             if cell_match and len(cell_match) >= 1:
                 self.num_lstm_layers = len(cell_match)
-                logger.debug("number of lstm layers: " + str(self.num_lstm_layers))
+                logger.debug("number of LSTM layers: %s", self.num_lstm_layers)
                 for i in range(self.num_lstm_layers):
                     self.state_variable_handlers.append({
                         "ct" + str(i): (self._ct_variable_finder, self._connect_lstm_yc_to_graph, i),
                         "ht" + str(i): (self._ht_variable_finder, self._connect_lstm_yh_to_graph, i)
                     })
                     self.state_variable_handlers.append({
-                         "ct_ht" + str(i): (self._ct_ht_shared_variable_finder, self._connect_lstm_ych_to_graph, i)
+                        "ct_ht" + str(i): (self._ct_ht_shared_variable_finder, self._connect_lstm_ych_to_graph, i)
                     })
                 logger.debug("parsing unit is %s, num layers is %d", cell_type, self.num_lstm_layers)
             if cell_match:
@@ -287,9 +287,9 @@ def process_var_init_nodes(self, context):
     def process_var_init_nodes_per_layer(self, context, i):
         init_h_id = None
         init_c_id = None
-        if ("ct_ht" + str(i)) in context.state_variables:
+        if "ct_ht" + str(i) in context.state_variables:
             init_h_id, init_c_id = self._process_non_tuple_ch_init_nodes(context, i)
-        elif ("ct" + str(i)) in context.state_variables and ("ht" + str(i)) in context.state_variables:
+        elif "ct" + str(i) in context.state_variables and ("ht" + str(i)) in context.state_variables:
             init_h_id, init_c_id = self._process_tuple_ch_init_nodes(context, i)
         else:
             raise ValueError("no initializers, unexpected")
@@ -363,11 +363,10 @@ def create_single_rnn_node(self, context, i):
     def create_rnn_node(self, context):
         rnn_nodes = list()
         outputs = context.loop_properties.scan_outputs_exits
-        logger.debug("number of rnn node outputs:" + str(len(outputs)))
-        for i in range(len(outputs)):
-            logger.debug("output " + str(i) + " with id=" + outputs[i].id)
+        logger.debug("number of rnn node outputs: %s", len(outputs))
+
         for i in range(self.num_lstm_layers):
-            logger.debug("creating rnn node for layer: " + str(i))
+            logger.debug("creating rnn node for layer: %s", i)
             rnn_nodes.append(self.create_single_rnn_node(context, i))
             output_id = rnn_nodes[i].output[0]
             rnn_output_shape = self.g.get_shape(output_id)
@@ -376,7 +375,7 @@ def create_rnn_node(self, context):
                                             shapes=[squeeze_output_shape],
                                             dtypes=[self.g.get_dtype(output_id)])
             if i + 1 < self.num_lstm_layers:
-                logger.debug("setting input for layer: " + str(i + 1))
+                logger.debug("setting input for layer: %s", i + 1)
                 context.onnx_input_ids[i + 1]["X"] = squeeze_node.output[0]
         return rnn_nodes
 
@@ -420,4 +419,4 @@ def _connect_lstm_ych_to_graph(self, context, i):
                                         shapes=[squeeze_output_shape],
                                         dtypes=[self.g.get_dtype(concat.output[0])])
 
-        self.g.replace_all_inputs(self.g.get_nodes(), exit_output.id, squeeze_node.output[0])
+        self.g.replace_all_inputs(self.g.get_nodes(), exit_output.id, squeeze_node.output[0])

tf2onnx/rewriter/lstm_rewriter_base.py

@@ -21,7 +21,7 @@
 logger = logging.getLogger(__name__)
 
 
-# pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test,broad-except,protected-access
+# pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test,broad-except,protected-access,W0223
 
 class LSTMContext(UnitRnnContext):
     def __init__(self):
@@ -45,6 +45,7 @@ class LSTMRewriterBase(UnitRnnRewriterBase):
         2 find needed info from tensorflow graph
     3 process found info according to ONNX requirement
     """
+
     def __init__(self, g):
         super(LSTMRewriterBase, self).__init__(g)
         # {var_name: (finder, connector)}
@@ -100,7 +101,8 @@ def parse_unit_rnn(self, context):
             context.input_size.append(None)
             context.hidden_size.append(None)
             context.attributes.append({})
-            context.onnx_input_ids[i]["sequence_lens"] = seq_len_node.output[0] if seq_len_node else utils.ONNX_EMPTY_INPUT
+            context.onnx_input_ids[i]["sequence_lens"] = \
+                seq_len_node.output[0] if seq_len_node else utils.ONNX_EMPTY_INPUT
 
         context.onnx_input_ids[0]["X"] = inputs[0]
         if not self.parse_attributes(context):
@@ -126,10 +128,9 @@ def _match_cell(self, context, unittype):
             )
 
             match_results = list(matcher.match_ops(body_graph_ops))
-            logger.debug("number of match results: " + str(len(match_results)))
-            if len(match_results) < 1:
-                return None
-            return match_results
+            logger.debug("number of match results: %s", len(match_results))
+            if len(match_results) > 0:
+                return match_results
         return None
 
     def get_state_variables(self, context):
@@ -186,4 +187,4 @@ def connect_unit_rnn_output_to_graph(self, context):
         squeeze_node = self.g.make_node("Squeeze", [output_id], attr={"axes": [1]},
                                         shapes=[squeeze_output_shape],
                                         dtypes=[self.g.get_dtype(output_id)])
-        self.g.replace_all_inputs(self.g.get_nodes(), gather_output_id, squeeze_node.output[0])
+        self.g.replace_all_inputs(self.g.get_nodes(), gather_output_id, squeeze_node.output[0])