support domain when make node

Author: nbcsm

tf2onnx/graph.py

@@ -378,7 +378,8 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
 
             new_outputs = [o if o != output else new_output_name for output in n.output]
             new_node = self.make_node(n.type, n.input, outputs=new_outputs, attr=n.attr, name=n.name,
-                                      skip_conversion=n._skip_conversion, dtypes=n_dtypes, shapes=n_shapes)
+                                      skip_conversion=n._skip_conversion, dtypes=n_dtypes, shapes=n_shapes,
+                                      domain=n.domain)
 
             if body_graphs:
                 for attr_name, body_graph in body_graphs.items():
@@ -423,7 +424,7 @@ def make_const(self, name, np_val, skip_conversion=False, raw=True):
         return node
 
     def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, skip_conversion=True,
-                  op_name_scope=None, name=None, shapes=None, dtypes=None):
+                  op_name_scope=None, name=None, shapes=None, dtypes=None, domain=None):
         """Make a new onnx node in the graph"""
         if attr is None:
             attr = {}
@@ -456,7 +457,7 @@ def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, sk
             n = self.get_node_by_output_in_current_graph(o)
             utils.make_sure(n is None, "output tensor named %s already exists in node: \n%s", o, n)
 
-        onnx_node = helper.make_node(op_type, inputs, outputs, name=name, **raw_attr)
+        onnx_node = helper.make_node(op_type, inputs, outputs, name=name, domain=domain, **raw_attr)
 
         if op_type in ["If", "Loop", "Scan"]:
             # we force the op containing inner graphs not skipped during conversion.
@@ -883,7 +884,7 @@ def remove_input(node, to_be_removed):
         # don't remove output from parent since others might depend on it
         return True
 
-    def insert_new_node_on_input(self, node, op_type, input_name, name=None, **kwargs):
+    def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=None, **kwargs):
         """Create and insert a new node into the graph.
         Args:
             node: we want to replace the input for this node
@@ -898,14 +899,14 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, **kwarg
         if name is None:
             name = utils.make_name(node.name)
         new_output = port_name(name)
-        new_node = self.make_node(op_type, [input_name], attr=kwargs, outputs=[new_output], name=name)
+        new_node = self.make_node(op_type, [input_name], attr=kwargs, outputs=[new_output], name=name, domain=domain)
         for i, n in enumerate(node.input):
             if n == input_name:
                 node.input[i] = new_output
                 break
         return new_node
 
-    def insert_new_node_on_output(self, op_type, output_name, name=None, **kwargs):
+    def insert_new_node_on_output(self, op_type, output_name, name=None, domain=None, **kwargs):
         """Create and insert a new node into the graph.
         Args:
             op_type: type for new operation
@@ -922,7 +923,7 @@ def insert_new_node_on_output(self, op_type, output_name, name=None, **kwargs):
                         type(op_type))
 
         new_output = port_name(name)
-        new_node = self.make_node(op_type, [output_name], attr=kwargs, outputs=[new_output], name=name)
+        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)