[export] Unflatten None (pytorch#153000)

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#153000
Approved by: https://github.com/pianpwk

Author: angelayi

test/export/test_unflatten.py

@@ -954,6 +954,27 @@ def forward(self, x):
         unflattened.foo = torch.compile(unflattened.foo, fullgraph=True)
         self.compare_outputs(orig_eager, unflattened, inputs)
 
+    def test_unflatten_none(self):
+        class M2(torch.nn.Module):
+            def forward(self, x, y):
+                return x + x, None
+
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.m2 = M2()
+
+            def forward(self, x, y):
+                x = x + x
+                return self.m2(x, y)
+
+        ep = export(
+            M(), (torch.rand(2, 3), None), preserve_module_call_signature=("m2",)
+        )
+        unflattened = unflatten(ep)
+        inp = (torch.randn(2, 3), None)
+        self.assertTrue(torch.allclose(M()(*inp)[0], unflattened(*inp)[0]))
+
 
 if __name__ == "__main__":
     run_tests()

torch/export/unflatten.py

@@ -1171,7 +1171,13 @@ def finalize_outputs(self):
             for output in signature.outputs:
                 if isinstance(
                     output,
-                    (TensorArgument, SymIntArgument, SymBoolArgument, SymFloatArgument),
+                    (
+                        TensorArgument,
+                        SymIntArgument,
+                        SymBoolArgument,
+                        SymFloatArgument,
+                        ConstantArgument,
+                    ),
                 ):
                     if output.name in self.seen_nodes:
                         orig_outputs.append(self.seen_nodes[output.name])