make ops optional in replace_all_inputs

Author: sdpython

tests/test_internals.py

@@ -139,7 +139,7 @@ def test_rewrite_subgraph(self):
             op_name = utils.make_name("ReplacedOp")
             out_name = utils.port_name(op_name)
             new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name)
-            g.replace_all_inputs(None, output_node.output[0], new_node.output[0])  # ops
+            g.replace_all_inputs(output_node.output[0], new_node.output[0])  # ops=ops
             for n in set(match.get_nodes()):
                 g.remove_node(n.name)
         g.topological_sort(ops)

tf2onnx/graph.py

@@ -473,7 +473,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(o, new_output_name, ops=self.get_nodes())
                 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)
@@ -839,7 +839,7 @@ def change_node_name(self, node, new_name):
                 if k == old_output:
                     self.outputs[j] = new_output
                     break
-            self.replace_all_inputs(self.get_nodes(), old_output, new_output)
+            self.replace_all_inputs(old_output, new_output, ops=self.get_nodes())
         return new_node
 
     def add_graph_input(self, name, dtype=None, shape=None):
@@ -1250,7 +1250,7 @@ def insert_new_node_on_output(self, op_type, output_name, name, domain=None, **k
 
         to_replace = [self.get_node_by_name(n) for n in self._input_to_node_name[output_name]]
         to_replace = [n for n in to_replace if n != new_node]
-        self.replace_all_inputs(to_replace, output_name, new_output)
+        self.replace_all_inputs(output_name, new_output, ops=to_replace)
         return new_node
 
     def find_output_consumers(self, output_name):
@@ -1298,11 +1298,11 @@ def _unregister_input_name(self, input_name, node, only_graph=False):
             del self.parent_graph._input_to_graph[input_name][id(self)]
             self.parent_graph._unregister_input_name(input_name, node, only_graph=True)
 
-    def replace_all_inputs(self, ops, old_input, new_input):
+    def replace_all_inputs(self, old_input, new_input, ops=None):
         """
         Replace all inputs pointing to old_input with new_input.
-        *ops* is used if defined, otherwise _input_to_node_name
-        is used to determine the impacted nodes.
+        *ops* is used if defined, otherwise `_input_to_node_name`
+        is used to determine the impacted nodes.        
         """
         if old_input == new_input:
             return
@@ -1333,7 +1333,8 @@ def replace_all_inputs(self, ops, old_input, new_input):
         # modify references in sub graphs
         if old_input in self._input_to_graph:
             for g in self._input_to_graph[old_input].values():
-                g.replace_all_inputs(g.get_nodes() if keep_ops else None, old_input, new_input)
+                g.replace_all_inputs(old_input, new_input,
+                                     ops=g.get_nodes() if keep_ops else None)
 
     def replace_input(self, node, old_input, new_input, input_index=None):
         """

tf2onnx/onnx_opset/controlflow.py

@@ -492,7 +492,7 @@ class TensorListStack:
     def version_7(cls, ctx, node, **kwargs):
         if node.inputs[0].is_while():
             ctx.remove_node(node.name)
-            ctx.replace_all_inputs(None, node.output[0], node.input[0])  # ctx.get_nodes()
+            ctx.replace_all_inputs(node.output[0], node.input[0])  # ops=ctx.get_nodes()
 
 
 @tf_op(["While", "StatelessWhile"])
@@ -582,7 +582,7 @@ def version_7(cls, ctx, node, **kwargs):
         for idx, n in reversed(to_remove):
             ctx.remove_node(n.name)
             # make the node output bad
-            ctx.replace_all_inputs(None, n.output[0], "@@ALLOC")  # ctx.get_nodes()
+            ctx.replace_all_inputs(n.output[0], "@@ALLOC")  # ops=ctx.get_nodes()
             del body.func_inputs[idx]
             del cond_graph.func_inputs[idx]
             del tf_while_inputs[idx]
@@ -618,7 +618,7 @@ def version_7(cls, ctx, node, **kwargs):
 
         # shift output consumers
         for k, v in output_map.items():
-            ctx.replace_all_inputs(None, k, v)  # ctx.get_nodes()
+            ctx.replace_all_inputs(k, v)  # ops=ctx.get_nodes()
 
         wire_while_body(ctx, body, loop_node.inputs, body_input_to_state_var, cond_input_to_state_var, output_shapes,
                         output_dtypes, body_name, node.name, cond_graph, tf_while_inputs, removed_scan_outputs)
@@ -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(old_output, new_output, ops=ops)
         to_remove.append(node)
     for node in to_remove:
         g.remove_node(node.name)

tf2onnx/onnx_opset/math.py

@@ -695,4 +695,4 @@ def atan2(y, x):
             "Add", inputs=[atan_node.output[0], pi_part.output[0]],
             op_name_scope=node.name + 'all',
             shapes=[shape], dtypes=[onnx_dtype])
-        ctx.replace_all_inputs(None, node.output[0], last_node.output[0])  # ctx.get_nodes()
+        ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()

tf2onnx/onnx_opset/misc.py

@@ -30,7 +30,7 @@ def version_1(cls, ctx, node, **kwargs):
             # if identity has a const as input, remove it
             input_name = node.input[0]
             output_name = node.output[0]
-            ctx.replace_all_inputs(None, output_name, input_name)  # ctx.get_nodes()
+            ctx.replace_all_inputs(output_name, input_name)  # ops=ctx.get_nodes()
             ctx.remove_node(node.name)
 
 

tf2onnx/onnx_opset/nn.py

@@ -451,7 +451,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(node.output[0], slice_node.output[0], ops=downstream_nodes)
 
         conv_dims_attr(node, "strides", spatial=spatial)
         conv_dims_attr(node, "dilations", spatial=spatial)

tf2onnx/onnx_opset/quantize.py

@@ -78,4 +78,4 @@ def version_10(cls, ctx, node, **kwargs):
             "DequantizeLinear", [new_node.output[0], pb_scale.name, zero_point.name],
             op_name_scope=node.name, attr={"axis": axis},
             shapes=[shape], dtypes=[dtype])
-        ctx.replace_all_inputs(None, node.output[0], last_node.output[0])  # ctx.get_nodes()
+        ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()

tf2onnx/onnx_opset/rnn.py

@@ -153,7 +153,7 @@ def make_sigmoid(i, w, b):
         h_node = ctx.make_node("Mul", [co_node.output[0], o])
 
         def replace_output(old_output, new_output):
-            ctx.replace_all_inputs(None, old_output, new_output)  # ctx.get_nodes()
+            ctx.replace_all_inputs(old_output, new_output)  # ops=ctx.get_nodes()
             ctx.copy_dtype(old_output, new_output)
             ctx.copy_shape(old_output, new_output)
 

tf2onnx/onnx_opset/tensor.py

@@ -115,7 +115,7 @@ def version_1(cls, ctx, node, **kwargs):
             # if identity has a const as input, remove it
             input_name = node.input[0]
             output_name = node.output[0]
-            ctx.replace_all_inputs(None, output_name, input_name)  # ctx.get_nodes()
+            ctx.replace_all_inputs(output_name, input_name)  # ops=ctx.get_nodes()
             ctx.remove_node(node.name)
 
 
@@ -125,7 +125,7 @@ class IdentityN:
     def version_1(cls, ctx, node, **kwargs):
         ctx.remove_node(node.name)
         for input_name, output_name in zip(node.input, node.output):
-            ctx.replace_all_inputs(None, output_name, input_name)  # ctx.get_nodes()
+            ctx.replace_all_inputs(output_name, input_name)  # ops=ctx.get_nodes()
 
 
 @tf_op("Reshape")
@@ -1050,7 +1050,7 @@ def version_1(cls, ctx, node, **kwargs):
         # concat all unqueezes
         concat = ctx.make_node("Concat", inputs, op_name_scope=node.name, attr={"axis": axis},
                                shapes=shapes, dtypes=dtypes)
-        ctx.replace_all_inputs(None, node.output[0], concat.output[0])  # ctx.get_nodes()
+        ctx.replace_all_inputs(node.output[0], concat.output[0])  # ops=ctx.get_nodes()
 
 
 @tf_op("Unpack")

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.output[0], node.input[0], ops=node2_consumers)
                 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.output[0], node.input[0], ops=node2_consumers)
         g.remove_node(node.name)
         g.remove_node(node2.name)
         return []