code refactor

Author: zhijxu-MS

tests/test_backend.py

@@ -721,9 +721,9 @@ def test_relu(self):
     @skip_caffe2_backend("fails on caffe2 with dim issue")
     @check_onnxruntime_incompatibility("Mul")
     def test_leaky_relu(self):
-        for alpha in [0.1, -0.1]:
-            x_val = np.array([0.5, 1.0, -0.5, -1.0], dtype=np.float32).reshape((2, 2))
-            x = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT)
+        for alpha in [0.1, -0.1, 1.0, -1.0, 10.0, -10.0]:
+            x_val = 1000*np.random.random_sample([1000, 100]).astype(np.float32)
+            x = tf.placeholder(tf.float32, [None]*x_val.ndim, name=_TFINPUT)
             x_ = tf.nn.leaky_relu(x, alpha)
             _ = tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val})

tf2onnx/graph.py

@@ -552,8 +552,7 @@ def _push_stack(stack, node, in_stack):
             stack.append(node)
             if node in in_stack:
                 raise ValueError('Graph has cycles.')
-            else:
-                in_stack[node] = True
+            in_stack[node] = True
 
         def _get_unvisited_child(g, node, not_visited):
             for child in g[node]:

tf2onnx/rewriter/leakyrelu_rewriter.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 """
-tf2onnx.rewrite - rewrite tensorflow subgraph to onnx leakyrelu op
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx leakyrelu op
 """
 
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
@@ -32,13 +32,13 @@ def rewrite_leakyrelu(g, ops):
         mul_node = match.get_op("mul")
         mul_input_node = match.get_op('mul_input')
 
-        max_input_edge_name = _find_edges_name_btw_nodes(max_input_node, max_node)
-        mul_input_edge_name = _find_edges_name_btw_nodes(mul_input_node, mul_node)
+        max_input_edge_name = _find_edge_name_between_nodes(max_input_node, max_node)
+        mul_input_edge_name = _find_edge_name_between_nodes(mul_input_node, mul_node)
         if max_input_edge_name == mul_input_edge_name:
             alpha = match.get_op("alpha").get_tensor_value()
             if alpha >= 1:
                 continue
-            leakyrelu = g.make_node("LeakyRelu", inputs=max_input_edge_name, attr={"alpha": alpha},
+            leakyrelu = g.make_node("LeakyRelu", inputs=[max_input_edge_name], attr={"alpha": alpha},
                                     shapes=[g.get_shape(max_node.output[0])], dtypes=[g.get_dtype(max_node.output[0])])
             ops.remove(max_node)
             ops.remove(mul_node)
@@ -48,10 +48,10 @@ def rewrite_leakyrelu(g, ops):
     return ops
 
 
-def _find_edges_name_btw_nodes(sender, sinker):
-    res = []
-    for sinker_end in sinker.input:
-        for sender_end in sender.output:
-            if sinker_end == sender_end:
-                res.append(sinker_end)
-    return res
+def _find_edge_name_between_nodes(src_node, consumer_node):
+    # find the first edge connection between two nodes.
+    for consumer_end in consumer_node.input:
+        for src_end in src_node.output:
+            if consumer_end == src_end:
+                return consumer_end
+    return None

tf2onnx/rewriter/random_uniform.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 """
-tf2onnx.rewrite - rewrite tensorflow subgraph to onnx random_uniform op
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx random_uniform op
 """
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx import utils

tf2onnx/shape_inference.py

@@ -98,8 +98,7 @@ def infer_shape_for_node(g, node):
             val = list(shape_attr.floats)
             if val:
                 raise ValueError("placeholder shape has floats value, and not scalar value")
-            else:
-                new_shape = ()
+            new_shape = ()
 
         if new_shape is not None:
             g.set_shape(node.output[0], new_shape)