Add TupleConstruct to allow (a, b)=SomeModule()

nni/retiarii/converter/graph_gen.py

@@ -315,6 +315,11 @@ def handle_single_node(node):
                 new_node = ir_graph.add_node(build_full_name(module_name, OpTypeName.ListConstruct, self.global_seq), node.kind())
                 node_index[node] = new_node
                 self._add_edge(ir_graph, node, graph_inputs, node_index, new_node, output_remap)
+            elif node.kind() == 'prim::TupleConstruct':
+                self.global_seq += 1
+                new_node = ir_graph.add_node(build_full_name(module_name, OpTypeName.TupleConstruct, self.global_seq), node.kind())
+                node_index[node] = new_node
+                self._add_edge(ir_graph, node, graph_inputs, node_index, new_node, output_remap)
             elif node.kind() == 'aten::append':
                 self.global_seq += 1
                 aten_node = ir_graph.add_node(build_full_name(module_name, BasicOpsPT[node.kind()], self.global_seq), node.kind())

nni/retiarii/converter/op_types.py

@@ -10,6 +10,7 @@ class OpTypeName(str, Enum):
     Attr = 'Attr'
     Constant = 'Constant'
     ListConstruct = 'ListConstruct'
+    TupleConstruct = 'TupleConstruct'
     LayerChoice = 'LayerChoice'
     InputChoice = 'InputChoice'
     ValueChoice = 'ValueChoice'

nni/retiarii/operation.py

@@ -121,6 +121,8 @@ def to_forward_code(self, field: str, output: str, inputs: List[str]) -> str:
             return f'{output} = {value}'
         elif self.type == 'prim::ListConstruct':
             return f'{output} = [{", ".join(inputs)}]'
+        elif self.type == 'prim::TupleConstruct':
+            return f'{output} = ({", ".join(inputs)})'
         elif self.type == 'prim::GetAttr':
             return f"{output} = {self.parameters['input']}.{self.parameters['name']}"
         elif self.type == 'aten::mean':