introduce keep_ops=True

sdpython · sdpython · commit 62d8b4bd56bd · 2020-08-10T17:45:12.000+02:00
diff --git a/tf2onnx/graph.py b/tf2onnx/graph.py
@@ -482,7 +482,7 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
                         body_graph.parent_graph = self
                         new_node.set_body_graph_as_attr(attr_name, body_graph)
 
-                self.replace_all_inputs(self.get_nodes(), o, new_output_name)
+                self.replace_all_inputs(self.get_nodes(), o, new_output_name, keep_ops=True)
                 self.make_node("Identity", [new_output_name], outputs=[o], op_name_scope=n.name + "_" + "graph_outputs")
                 self.copy_shape(new_output_name, o)
                 self.copy_dtype(new_output_name, o)
@@ -1269,7 +1269,7 @@ def insert_new_node_on_output(self, op_type, output_name, name, domain=None, **k
         new_node = self.make_node(op_type, [output_name], attr=kwargs, outputs=[new_output], name=name, domain=domain)
 
         to_replace = [n for n in self.get_nodes() if n != new_node]
-        self.replace_all_inputs(to_replace, output_name, new_output)
+        self.replace_all_inputs(to_replace, output_name, new_output, keep_ops=True)
         return new_node
 
     def find_output_consumers(self, output_name):
@@ -1286,8 +1286,11 @@ def find_output_consumers(self, output_name):
                     nodes.extend(g.find_output_consumers(output_name))
         return nodes
 
-    def replace_all_inputs(self, ops, old_input, new_input):
-        """Replace all inputs pointing to old_input with new_input."""
+    def replace_all_inputs(self, ops, old_input, new_input, keep_ops=True):
+        """
+        Replace all inputs pointing to old_input with new_input.
+        *ops* is unused unless keep_ops is True.
+        """
         if old_input == new_input:
             return
         if new_input not in self._input_to_node_name:
diff --git a/tf2onnx/onnx_opset/controlflow.py b/tf2onnx/onnx_opset/controlflow.py
@@ -813,7 +813,7 @@ def prefix_graph(g, scope):
                 if old_output == oname:
                     g.outputs[i] = new_output
                     break
-            g.replace_all_inputs(ops, old_output, new_output)
+            g.replace_all_inputs(ops, old_output, new_output, keep_ops=True)
         to_remove.append(node)
     for node in to_remove:
         g.remove_node(node.name)
diff --git a/tf2onnx/onnx_opset/nn.py b/tf2onnx/onnx_opset/nn.py
@@ -447,7 +447,7 @@ def version_1(cls, ctx, node, **kwargs):
             downstream_nodes = ctx.find_output_consumers(node.output[0])
             downstream_nodes.remove(output_shape)
             downstream_nodes.remove(slice_node)
-            ctx.replace_all_inputs(downstream_nodes, node.output[0], slice_node.output[0])
+            ctx.replace_all_inputs(downstream_nodes, node.output[0], slice_node.output[0], keep_ops=True)
 
         conv_dims_attr(node, "strides", spatial=spatial)
         conv_dims_attr(node, "dilations", spatial=spatial)
diff --git a/tf2onnx/onnx_opset/tensor.py b/tf2onnx/onnx_opset/tensor.py
@@ -272,7 +272,7 @@ def version_1(cls, ctx, node, **kwargs):
         axis_node = node.inputs[-1]
         utils.make_sure(axis_node.is_const(), "%r needs to be const", axis_node.name)
         axis_val = axis_node.get_tensor_value()
-        ctx.remove_input(node, node.input[-1])
+        ctx.remove_input(node, node.input[-1], len(node.input) - 1)
 
         if axis_val < 0:  # onnxruntime does not support -1 axis, but TF supports.
             input_shape = ctx.get_shape(node.input[0])
diff --git a/tf2onnx/optimizer/back_to_back_optimizer.py b/tf2onnx/optimizer/back_to_back_optimizer.py
@@ -138,7 +138,7 @@ def _optimize_transpose(g, node, consumer_nodes):
                 shape = g.get_shape(node2.output[0])
                 dtype = g.get_dtype(node2.output[0])
                 node2_consumers = g.find_output_consumers(node2.output[0])
-                g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0])
+                g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0], keep_ops=True)
                 g.remove_node(node2.name)
                 if set(node2.output) & set(g.outputs):
                     g.make_node("Identity", [node.input[0]],
@@ -173,7 +173,7 @@ def _optimize_squeeze_unsqueeze(g, node, consumer_nodes):
             return []
 
         node2_consumers = g.find_output_consumers(node2.output[0])
-        g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0])
+        g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0], keep_ops=True)
         g.remove_node(node.name)
         g.remove_node(node2.name)
         return []
diff --git a/tf2onnx/optimizer/const_fold_optimizer.py b/tf2onnx/optimizer/const_fold_optimizer.py
@@ -90,7 +90,7 @@ def _replace_node_with_const(node, graph, vals):
             const_node = graph.make_const(utils.make_name("const_fold_opt"), val)
             graph.set_dtype(const_node.output[0], utils.map_numpy_to_onnx_dtype(val.dtype))
             graph.set_shape(const_node.output[0], val.shape)
-            graph.replace_all_inputs(graph.get_nodes(), old_input, const_node.output[0])
+            graph.replace_all_inputs(graph.get_nodes(), old_input, const_node.output[0], keep_ops=True)
         graph.remove_node(node.name)
 
     @staticmethod
diff --git a/tf2onnx/optimizer/identity_optimizer.py b/tf2onnx/optimizer/identity_optimizer.py
@@ -47,7 +47,7 @@ def _optimize_at_current_graph_level(self, g):
     def _handle_non_graph_output_identity(graph, identity):
         old_name = identity.output[0]
         new_name = identity.input[0]
-        graph.replace_all_inputs(graph.get_nodes(), old_name, new_name)
+        graph.replace_all_inputs(graph.get_nodes(), old_name, new_name, keep_ops=True)
         graph.remove_node(identity.name)
         return True
 
@@ -81,5 +81,5 @@ def _handle_graph_output_identity(self, graph, identity, graph_outputs):
         graph.set_shape(output_id, output_shape)
         graph.set_dtype(output_id, output_dtype)
 
-        graph.replace_all_inputs(graph.get_nodes(), input_id, output_id)
+        graph.replace_all_inputs(graph.get_nodes(), input_id, output_id, keep_ops=True)
         return True
diff --git a/tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py b/tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py
@@ -99,7 +99,7 @@ def _merge_nodes_that_are_duplicated(self, nodes_to_process, graph):
             if set(node_to_delete.output).intersection(set(graph.outputs)):
                 continue
             for old_input, new_input in zip(node_to_delete.output, node_to_retain.output):
-                graph.replace_all_inputs(graph.get_nodes(), old_input, new_input)
+                graph.replace_all_inputs(graph.get_nodes(), old_input, new_input, keep_ops=True)
             graph.remove_node(node_to_delete.name)
             self._graph_can_be_optimized = True
 
diff --git a/tf2onnx/optimizer/transpose_optimizer.py b/tf2onnx/optimizer/transpose_optimizer.py
@@ -408,9 +408,8 @@ def _add_handler(self, trans, node):
 
                 conv_inputs = [t_p.input[0], t_p.input[1], node.input[1]]
                 conv_node = self._g.make_node(t_p.type, conv_inputs, attr=t_p.attr_onnx)
-                ops = self._g.get_nodes()
                 self._g.replace_input(trans, trans.input[0], utils.port_name(conv_node.name), 0)
-                self._g.replace_all_inputs(ops, node.output[0], trans.output[0])
+                self._g.replace_all_inputs(self._g.get_nodes(), node.output[0], trans.output[0])
                 self._g.remove_node(t_p.name)
                 self._g.remove_node(node.name)
                 return True
@@ -419,8 +418,7 @@ def _add_handler(self, trans, node):
     def _transpose_handler(self, trans, node):
         if is_nchw_transpose(node):
             for g in {self._g, node.graph}:
-                ops = g.get_nodes()
-                g.replace_all_inputs(ops, node.output[0], trans.input[0])
+                g.replace_all_inputs(g.get_nodes(), node.output[0], trans.input[0])
 
             shape = node.graph.get_shape(node.output[0])
             dtype = node.graph.get_dtype(node.output[0])
diff --git a/tf2onnx/rewriter/custom_rnn_rewriter.py b/tf2onnx/rewriter/custom_rnn_rewriter.py
@@ -151,19 +151,23 @@ def _connect_scan_with_output(self, context, scan_node):
                 if self.g.opset == 8:
                     nodes = self._adapt_scan_sequence_input_or_output("state_output_reshape",
                                                                       scan_node.output[index], True)
-                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
+                    self.g.replace_all_inputs(
+                        self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
                 else:  # since opset 9
-                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, scan_node.output[index])
+                    self.g.replace_all_inputs(
+                        self.g.get_nodes(), out_tensor_value_info.id, scan_node.output[index])
             index += 1
 
         for out_tensor_value_info in context.loop_properties.scan_outputs_exits:
             if out_tensor_value_info.id:
                 if self.g.opset == 8:
                     nodes = self._adapt_scan_sequence_input_or_output("scan_output_reshape",
                                                                       scan_node.output[index], True)
-                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
+                    self.g.replace_all_inputs(
+                        self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
                 else:  # since opset 9
-                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, scan_node.output[index])
+                    self.g.replace_all_inputs(
+                        self.g.get_nodes(), out_tensor_value_info.id, scan_node.output[index])
             index += 1
 
     def _adapt_scan_sequence_input_or_output(self, target_name, input_id, handle_output=False):
diff --git a/tf2onnx/rewriter/dropout_rewriter.py b/tf2onnx/rewriter/dropout_rewriter.py
@@ -96,7 +96,7 @@ def rewrite_dropout(g, ops):
                 shapes=[g.get_shape(data.output[0])],
                 dtypes=[g.get_dtype(data.output[0])]
             )
-            g.replace_all_inputs(ops, outputs.output[0], new_node.output[0])
+            g.replace_all_inputs(ops, outputs.output[0], new_node.output[0], keep_ops=True)
             for n in nodes_to_remove:
                 g.remove_node(n.name)
 
diff --git a/tf2onnx/rewriter/flatten_rewriter.py b/tf2onnx/rewriter/flatten_rewriter.py
@@ -94,7 +94,7 @@ def rewrite_flatten(g, ops):
                 new_dim = -1
 
             g.set_shape(out_name, input_shape[:-2] + [new_dim])
-            g.replace_all_inputs(ops, reshape_node.output[0], out_name)
+            g.replace_all_inputs(ops, reshape_node.output[0], out_name, keep_ops=True)
             for n in to_remove:
                 g.remove_node(n.name)
 
diff --git a/tf2onnx/rewriter/gemm_rewriter.py b/tf2onnx/rewriter/gemm_rewriter.py
@@ -95,7 +95,7 @@ def rewrite_gemm(g, ops):
                                    dtypes=[g.get_dtype(add_node.output[0])], op_name_scope=matmul_node.name)
 
                 ops.append(gemm)
-                g.replace_all_inputs(ops, add_node.output[0], gemm.output[0])
+                g.replace_all_inputs(ops, add_node.output[0], gemm.output[0], keep_ops=True)
                 to_delete = [add_node, matmul_node]
                 g.safe_remove_nodes(to_delete)
     return ops
diff --git a/tf2onnx/rewriter/gru_rewriter.py b/tf2onnx/rewriter/gru_rewriter.py
@@ -212,7 +212,7 @@ def process_var_init_nodes(self, context):
             return
         squeeze_node = self.g.make_node("Unsqueeze", [initializer_input_id], attr={"axes": [0]})
         to_replace = [n for n in self.g.get_nodes() if n != squeeze_node]
-        self.g.replace_all_inputs(to_replace, initializer_input_id, squeeze_node.output[0])
+        self.g.replace_all_inputs(to_replace, initializer_input_id, squeeze_node.output[0], keep_ops=True)
         context.onnx_input_ids["initial_state"] = squeeze_node.output[0]
 
     def create_rnn_node(self, context):
diff --git a/tf2onnx/rewriter/leakyrelu_rewriter.py b/tf2onnx/rewriter/leakyrelu_rewriter.py
@@ -41,7 +41,7 @@ def rewrite_leakyrelu(g, ops):
             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.append(leakyrelu)
-            g.replace_all_inputs(ops, max_node.output[0], leakyrelu.output[0])
+            g.replace_all_inputs(ops, max_node.output[0], leakyrelu.output[0], keep_ops=True)
             to_delete = [max_node, mul_node]
             g.safe_remove_nodes(to_delete)
 
diff --git a/tf2onnx/rewriter/loop_rewriter.py b/tf2onnx/rewriter/loop_rewriter.py
@@ -85,7 +85,8 @@ def rewrite(self, context):
                 index_node = loop_body_g.make_node("Unsqueeze", [input_ta.index_input_id], attr={"axes": [0]})
                 gather_node = loop_body_g.make_node("Gather", [input_ta.data_input_id, index_node.output[0]])
                 data_node = loop_body_g.make_node("Squeeze", [gather_node.output[0]], attr={"axes": [0]})
-                loop_body_g.replace_all_inputs(loop_body_g.get_nodes(), input_ta.consumer.id, data_node.output[0])
+                loop_body_g.replace_all_inputs(
+                    loop_body_g.get_nodes(), input_ta.consumer.id, data_node.output[0])
 
             ## create Loop node
             loop_node = self._create_loop_node(context, loop_props, init_cond_output)
@@ -134,10 +135,8 @@ def _create_subgraph_initial_cond(self, cond_graph):
         # replace all inputs of condition graph by initializer (enter_input)
         for loop_var in cond_graph.dependent_vars:
             self.g.replace_all_inputs(
-                copied_nodes,
-                loop_var.next_iteration_input.id,
-                loop_var.enter_input_id
-            )
+                copied_nodes, loop_var.next_iteration_input.id,
+                loop_var.enter_input_id, keep_ops=True)
         init_cond_output = "{}/{}".format(name_scope, cond_graph.outputs[0].id)
         self.g.set_dtype(init_cond_output, cond_graph.outputs[0].dtype)
         self.g.set_shape(init_cond_output, cond_graph.outputs[0].shape)
diff --git a/tf2onnx/rewriter/loop_rewriter_base.py b/tf2onnx/rewriter/loop_rewriter_base.py
@@ -275,7 +275,7 @@ def _crop_loop_body_sub_graph(self, context):
 
         for enter_node in enter_nodes:
             # connect Enter's output to Enter's input
-            self.g.replace_all_inputs(ops, enter_node.output[0], enter_node.input[0])
+            self.g.replace_all_inputs(ops, enter_node.output[0], enter_node.input[0], keep_ops=True)
 
         return GraphInfo(ops, inputs, outputs)
 
@@ -287,7 +287,7 @@ def _crop_loop_condition_sub_graph(self, context):
 
         for enter_node in enter_nodes:
             # connect Enter's output to Enter's input
-            self.g.replace_all_inputs(ops, enter_node.output[0], enter_node.input[0])
+            self.g.replace_all_inputs(ops, enter_node.output[0], enter_node.input[0], keep_ops=True)
 
         dependent_vars = []
         for merge_node in merge_nodes:
@@ -299,11 +299,11 @@ def _crop_loop_condition_sub_graph(self, context):
             # to consumer cell graph outputs.
             non_switch_consumers = [n for n in self.g.find_output_consumers(merge_node.output[0]) if n.type != "Switch"]
             self.g.replace_all_inputs(non_switch_consumers, merge_node.output[0],
-                                      loop_var.next_iteration_input.id)
+                                      loop_var.next_iteration_input.id, keep_ops=True)
             dependent_vars.append(loop_var)
 
         # cut off connection between condition graph and LoopCond node.
-        self.g.replace_all_inputs([context.loop_cond], context.loop_cond.output[0], INVALID_INPUT_ID)
+        self.g.replace_all_inputs([context.loop_cond], context.loop_cond.output[0], INVALID_INPUT_ID, keep_ops=True)
 
         graph_info = GraphInfo(ops, [], outputs)
         graph_info.dependent_vars = dependent_vars
@@ -319,11 +319,11 @@ 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_tf_tensor_array_write_op(n)]
-                self.g.replace_all_inputs(ta_write_nodes, val.next_iteration_input.id, INVALID_INPUT_ID)
+                self.g.replace_all_inputs(ta_write_nodes, val.next_iteration_input.id, INVALID_INPUT_ID, keep_ops=True)
             else:
                 # connect NextIteration to an invalid node, to cut off an ending node of the cell.
                 next_iter_nodes = [n for n in self.g.get_nodes() if n.type == "NextIteration"]
-                self.g.replace_all_inputs(next_iter_nodes, val.next_iteration_input.id, INVALID_INPUT_ID)
+                self.g.replace_all_inputs(next_iter_nodes, val.next_iteration_input.id, INVALID_INPUT_ID, keep_ops=True)
 
         for scan_input in context.loop_properties.scan_inputs:
             # remove the node to cut off connection between scan_input and the cell.
diff --git a/tf2onnx/rewriter/random_normal_rewriter.py b/tf2onnx/rewriter/random_normal_rewriter.py
@@ -55,6 +55,6 @@ def rewrite_random_normal(g, ops):
                 new_node = g.make_node("RandomNormal", [], outputs=[out_name], name=op_name,
                                        attr={"shape": shape, "mean": mean, "scale": 1.0, "dtype": dtype, "seed": seed})
 
-            g.replace_all_inputs(ops, output.output[0], new_node.output[0])
+            g.replace_all_inputs(ops, output.output[0], new_node.output[0], keep_ops=True)
             g.safe_remove_nodes(match.get_nodes())
     return ops
diff --git a/tf2onnx/rewriter/random_uniform.py b/tf2onnx/rewriter/random_uniform.py
@@ -32,7 +32,7 @@ def rewrite_random_uniform(g, ops):
         tmin = input2.inputs[1].get_tensor_value()
         to_delete = list(set(match.get_nodes()))
         new_node = create_onnx_random_uniform_op(g, tmax, tmin, ru_op, output, to_delete)
-        g.replace_all_inputs(ops, output.output[0], new_node.output[0])
+        g.replace_all_inputs(ops, output.output[0], new_node.output[0], keep_ops=True)
         g.safe_remove_nodes(to_delete)
 
     return ops
@@ -61,7 +61,7 @@ def rewrite_random_uniform_fold_const(g, ops):
         tmax = tmin + tmax_minus_tmin
         to_delete = list(set(match.get_nodes()))
         new_node = create_onnx_random_uniform_op(g, tmax, tmin, ru_op, output, to_delete)
-        g.replace_all_inputs(ops, output.output[0], new_node.output[0])
+        g.replace_all_inputs(ops, output.output[0], new_node.output[0], keep_ops=True)
         g.safe_remove_nodes(to_delete)
 
     return ops
diff --git a/tf2onnx/rewriter/rnn_utils.py b/tf2onnx/rewriter/rnn_utils.py
@@ -545,7 +545,7 @@ def slice_birnn_for_original_rnn_consumers(g, rnn_fw, rnn_bw, bi_rnn, rnn_output
 
         for r_op in reverse_nodes:
             logger.debug("remove reverse op %s", r_op.name)
-            g.replace_all_inputs(all_nodes, r_op.output[0], r_op.input[0])
+            g.replace_all_inputs(all_nodes, r_op.output[0], r_op.input[0], keep_ops=True)
             to_remove.append(r_op.name)
     elif rnn_output_index in [1, 2]:
         axis = 0
@@ -557,14 +557,14 @@ def slice_birnn_for_original_rnn_consumers(g, rnn_fw, rnn_bw, bi_rnn, rnn_output
         inputs_map = {"data": bi_rnn.output[rnn_output_index], **attr}
         slice_node_fw = GraphBuilder(g).make_slice(inputs_map)
         all_nodes.append(g.get_node_by_output(slice_node_fw))
-        g.replace_all_inputs(fw_consumers, rnn_fw.output[rnn_output_index], slice_node_fw)
+        g.replace_all_inputs(fw_consumers, rnn_fw.output[rnn_output_index], slice_node_fw, keep_ops=True)
 
     if bw_consumers:
         attr = {"axes": [axis], "starts": [1], "ends": [2]}
         inputs_map = {"data": bi_rnn.output[rnn_output_index], **attr}
         slice_node_bw = GraphBuilder(g).make_slice(inputs_map)
         all_nodes.append(g.get_node_by_output(slice_node_bw))
-        g.replace_all_inputs(bw_consumers, rnn_bw.output[rnn_output_index], slice_node_bw)
+        g.replace_all_inputs(bw_consumers, rnn_bw.output[rnn_output_index], slice_node_bw, keep_ops=True)
 
 
 def remove_reverse_in_bw_input(g, bw_rnn_input_x, rnn_type):
diff --git a/tf2onnx/rewriter/thresholded_relu_rewriter.py b/tf2onnx/rewriter/thresholded_relu_rewriter.py
@@ -43,7 +43,7 @@ def rewrite_thresholded_relu(g, ops):
             thresholded_relu = g.make_node("ThresholdedRelu", inputs=[mul_input_edge_name], attr={"alpha": theta},
                                            shapes=[g.get_shape(mul_node.output[0])],
                                            dtypes=[g.get_dtype(mul_node.output[0])])
-            g.replace_all_inputs(ops, mul_node.output[0], thresholded_relu.output[0])
+            g.replace_all_inputs(ops, mul_node.output[0], thresholded_relu.output[0], keep_ops=True)
             to_delete = [cast_node, mul_node]
             g.safe_remove_nodes(to_delete)
     return ops
