Handle unpacking during TorchScript to ExportedProgram conversion (#127419)

jiashenC · pytorchmergebot · commit 4f9fcd71562a · 2024-06-05T15:27:13.000Z
Pull Request resolved: #127419 Approved by: https://github.com/angelayi
diff --git a/test/export/test_converter.py b/test/export/test_converter.py
@@ -247,6 +247,23 @@ def forward(
         inp = (torch.randn(10, 10), torch.rand(10, 10))
         self._check_equal_ts_ep_converter(Module(), inp)
 
+    def test_ts2ep_converter_unpack(self):
+        class MUnpackList(torch.nn.Module):
+            def forward(self, x):
+                x, y = torch.split(x, 2)
+                return x + y
+
+        class MUnpackTuple(torch.nn.Module):
+            def forward(self, x_tuple: Tuple[torch.Tensor, torch.Tensor]):
+                x, y = x_tuple
+                x = x.cos()
+                return x + y
+
+        inp = torch.ones(1, 4)
+        self._check_equal_ts_ep_converter(MUnpackList(), inp)
+        inp = ((torch.zeros(1, 4), torch.ones(1, 4)),)
+        self._check_equal_ts_ep_converter(MUnpackTuple(), inp)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/torch/_export/converter.py b/torch/_export/converter.py
@@ -319,6 +319,20 @@ def convert_prim_DictConstruct(self, node: torch._C.Node):
         output_name = node.output().debugName()
         self.name_to_node[output_name] = output_dict
 
+    def convert_prim_ListUnpack(self, node: torch._C.Node):
+        self._convert_prim_unpack_iterator(node)
+
+    def convert_prim_TupleUnpack(self, node: torch._C.Node):
+        self._convert_prim_unpack_iterator(node)
+
+    def _convert_prim_unpack_iterator(self, node: torch._C.Node):
+        # Single input and multiple outputs for unpacking.
+        for i, outp in enumerate(node.outputs()):
+            outp_name = outp.debugName()
+            inp = self.get_fx_value(node.input())
+            fx_node = self.fx_graph.call_function(operator.getitem, (inp, i))
+            self.name_to_node[outp_name] = fx_node
+
     def convert_aten_Int(self, node: torch._C.Node):
         # converts aten::Int as aten._to_copy + aten::_local_scalar_dense
         target = torch.ops.aten._to_copy.default
