Skip to content

Commit c81962f

Browse files
angelayifacebook-github-bot
authored andcommitted
Serialize debug handle
Summary: Unfortunately we can't serialize the debug handle for placeholder/output nodes for now because they aren't stored as nodes... Reviewed By: zhxchen17 Differential Revision: D47311042 fbshipit-source-id: 296e3cb213160319cb37cbf6285f216b36f16c07
1 parent a55d6b6 commit c81962f

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

exir/serde/serialize.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,15 @@ def handle_call_function(self, node: torch.fx.Node) -> None:
4141

4242
super().handle_call_function(node)
4343

44+
def serialize_metadata(self, node: torch.fx.Node) -> Dict[str, str]:
45+
meta = super().serialize_metadata(node)
46+
47+
if "debug_handle" in node.meta:
48+
debug_handle = node.meta["debug_handle"]
49+
meta["debug_handle"] = str(debug_handle)
50+
51+
return meta
52+
4453
def serialize_alloc_inputs(
4554
self, inputs # pyre-ignore
4655
) -> List[schema.NamedArgument]:
@@ -198,6 +207,14 @@ def fake_op(x):
198207

199208
super().deserialize_node(serialized_node, target)
200209

210+
def deserialize_metadata(self, metadata: Dict[str, str]) -> Dict[str, Any]:
211+
res = super().deserialize_metadata(metadata)
212+
213+
if debug_handle := metadata.get("debug_handle"):
214+
res["debug_handle"] = int(debug_handle)
215+
216+
return res
217+
201218
# pyre-ignore
202219
def deserialize_alloc_inputs(self, serialized_inputs: List[schema.NamedArgument]):
203220
def deserialize_alloc_spec(serialized_alloc_spec: str) -> memory.AllocSpec:

0 commit comments

Comments
 (0)