(self, node: torch.fx.Node)
| 572 | pass |
| 573 | |
| 574 | def serialize_metadata(self, node: torch.fx.Node) -> Dict[str, str]: |
| 575 | ret = {} |
| 576 | if stack_trace := node.meta.get("stack_trace"): |
| 577 | ret["stack_trace"] = stack_trace |
| 578 | |
| 579 | if nn_module_stack := node.meta.get("nn_module_stack"): |
| 580 | |
| 581 | def export_nn_module_stack(val): |
| 582 | assert isinstance(val, tuple) and len(val) == 2 |
| 583 | path, ty = val |
| 584 | |
| 585 | assert isinstance(path, str) |
| 586 | |
| 587 | # node.meta["nn_module_stack"] could have two forms: |
| 588 | # 1. (path: str, module_type: 'type'), e.g. |
| 589 | # ('', <class 'sigmoid.inference.MySimpleModel'>) |
| 590 | # 2. (path: str, module_type: str), e.g. |
| 591 | # ('', 'sigmoid.inference.MySimpleModel') |
| 592 | # ExportedProgram directly produced by torch.export() has form 1 |
| 593 | # ExportedProgram deserialized from disk has form 2 |
| 594 | # TODO: This is not ideal, we should fix this. |
| 595 | if isinstance(ty, str): |
| 596 | normalized_ty = ty |
| 597 | else: |
| 598 | normalized_ty = ty.__module__ + "." + ty.__qualname__ |
| 599 | |
| 600 | return path + "," + normalized_ty |
| 601 | |
| 602 | # Serialize to "key,orig_path,type_str" |
| 603 | nn_module_list = [ |
| 604 | f"{k},{export_nn_module_stack(v)}" for k, v in nn_module_stack.items() |
| 605 | ] |
| 606 | ret["nn_module_stack"] = ST_DELIMITER.join(nn_module_list) |
| 607 | |
| 608 | if source_fn_st := node.meta.get("source_fn_stack"): |
| 609 | source_fn_list = [ |
| 610 | f"{source_fn[0]},{self.serialize_operator(source_fn[1])}" |
| 611 | for source_fn in source_fn_st |
| 612 | ] |
| 613 | ret["source_fn_stack"] = ST_DELIMITER.join(source_fn_list) |
| 614 | |
| 615 | if torch_fn := node.meta.get("torch_fn"): |
| 616 | ret["torch_fn"] = ST_DELIMITER.join(list(torch_fn)) |
| 617 | |
| 618 | return ret |
| 619 | |
| 620 | def serialize_script_obj_meta( |
| 621 | self, script_obj_meta: ep.CustomObjArgument |
no test coverage detected