[export] Fix serialize nn_module_stack (#104721)

angelayi · pytorchmergebot · commit 29c30b1db812 · 2023-07-07T17:13:17.000Z
Summary: Some serialized nn_module_stacks contain nested commas, something like: `(getitem(L['module'],0),torch.nn.modules.linear.Linear)` Fixing the parsing so that we can deserialize the string in the format of: `(local identifier, module type)` Test Plan: CI Differential Revision: D47252881 Pull Request resolved: #104721 Approved by: https://github.com/zhxchen17
diff --git a/torch/_export/serde/serialize.py b/torch/_export/serde/serialize.py
@@ -236,10 +236,21 @@ def deserialize_meta_func(serialized_target: str):
             key = kv[:key_idx]
             assert kv[key_idx + 1] == "("
             assert kv[-1] == ")"
-            values = kv[key_idx + 2: -1].split(",")
-            assert len(values) == 2
-            module = deserialize_meta_func(values[1])
-            nn_module_stack[key] = (values[0], module)
+
+            paren = 0
+            comma_idx = None
+            for i, c in enumerate(kv[key_idx + 2:-1]):
+                if c == "," and paren == 0:
+                    assert comma_idx is None
+                    comma_idx = i + key_idx + 2
+                elif c == "(":
+                    paren += 1
+                elif c == ")":
+                    paren -= 1
+
+            assert comma_idx is not None
+            module = deserialize_meta_func(kv[comma_idx + 1:-1])
+            nn_module_stack[key] = (kv[key_idx + 2:comma_idx], module)
         ret["nn_module_stack"] = nn_module_stack
 
     if source_fn_str := metadata.get("source_fn"):
