Load missing state dict in edge program serialization (#3076)

Summary:

The state dict wasn't being passed in when ExportedProgram was being created after deserialization.

Reviewed By: pssrawat

Differential Revision: D56224054

Author: tarun292

exir/program/_program.py

@@ -743,7 +743,7 @@ class EdgeProgramManager:
 
     def __init__(
         self,
-        edge_programs: Dict[str, ExportedProgram],
+        edge_programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
         constant_methods: Optional[Dict[str, Any]] = None,
         compile_config: Optional[EdgeCompileConfig] = None,
     ):
@@ -753,6 +753,8 @@ def __init__(
         Constructs an EdgeProgramManager from an existing set of exported programs in edge dialect.
         """
         config = compile_config or EdgeCompileConfig()
+        if not isinstance(edge_programs, dict):
+            edge_programs = {"forward": edge_programs}
         for name, program in edge_programs.items():
             try:
                 EXIREdgeDialectVerifier(

exir/serde/serialize.py

@@ -680,7 +680,7 @@ def deserialize(
             root=state_dict,
             graph=dummy_g,
             graph_signature=ep.ExportGraphSignature(input_specs=[], output_specs=[]),
-            state_dict={},  # TODO(T157676982)
+            state_dict=state_dict,  # TODO(T157676982)
             range_constraints=range_constraints,
             module_call_graph=module_call_graph,
             verifier=load_verifier(

exir/tests/test_serde.py

@@ -159,6 +159,28 @@ def forward(self, x):
         edge_new = deserialize(serialize(edge.exported_program()))
         self.check_ep(edge.exported_program(), edge_new, model_inputs)
 
+    def test_model_with_weights(self) -> None:
+        class LinearAdd(nn.Module):
+            def __init__(self, M: int, N: int):
+                super().__init__()
+                self.M = M
+                self.N = N
+                self.linear = torch.nn.Linear(M, N)
+
+            def forward(self, x, y):
+                x = self.linear(x)
+                y = self.linear(y)
+                return torch.add(x, y)
+
+            @classmethod
+            def _get_random_inputs(cls):
+                return (torch.rand(128, 20), torch.rand(128, 20))
+
+        linear_add = LinearAdd(20, 30)
+        model_inputs = LinearAdd._get_random_inputs()
+
+        self.check_serde(linear_add, model_inputs)
+
     def test_delegate_partitioner(self) -> None:
         class Model(torch.nn.Module):
             def __init__(self):