| 433 | |
| 434 | |
| 435 | class GraphModuleSerializer: |
| 436 | def __init__( |
| 437 | self, |
| 438 | graph_signature: ep.ExportGraphSignature, |
| 439 | module_call_graph: List[ep.ModuleCallEntry], |
| 440 | ): |
| 441 | self.graph_state = GraphState() |
| 442 | self.graph_signature = graph_signature |
| 443 | self.module_call_graph = module_call_graph |
| 444 | self.custom_objs: Dict[str, torch._C.ScriptObject] = {} |
| 445 | |
| 446 | @contextmanager |
| 447 | def save_graph_state(self): |
| 448 | saved = self.graph_state |
| 449 | self.graph_state = GraphState() |
| 450 | try: |
| 451 | yield |
| 452 | finally: |
| 453 | self.graph_state = saved |
| 454 | |
| 455 | def handle_placeholder(self, node: torch.fx.Node): |
| 456 | assert node.op == "placeholder" |
| 457 | if isinstance(node.meta["val"], torch.Tensor): |
| 458 | graph_input = Argument.create(as_tensor=TensorArgument(name=node.name)) |
| 459 | self.graph_state.tensor_values[node.name] = serialize_tensor_meta( |
| 460 | node.meta["val"] |
| 461 | ) |
| 462 | elif isinstance(node.meta["val"], torch.SymInt): |
| 463 | graph_input = Argument.create( |
| 464 | as_sym_int=SymIntArgument.create(as_name=node.name) |
| 465 | ) |
| 466 | self.graph_state.sym_int_values[node.name] = serialize_sym_int( |
| 467 | node.meta["val"] |
| 468 | ) |
| 469 | elif isinstance(node.meta["val"], (int, bool, str, float, type(None))): |
| 470 | graph_input = self.serialize_input(node.meta["val"]) |
| 471 | elif isinstance(node.meta["val"], ep.CustomObjArgument): |
| 472 | class_fqn = node.meta["val"].class_fqn |
| 473 | graph_input = Argument.create( |
| 474 | as_custom_obj=CustomObjArgument(name=node.name, class_fqn=class_fqn) |
| 475 | ) |
| 476 | self.graph_state.custom_obj_values[node.name] = ( |
| 477 | self.serialize_script_obj_meta(node.meta["val"]) |
| 478 | ) |
| 479 | else: |
| 480 | raise AssertionError(f"Unimplemented graph input type: {node.meta['val']}") |
| 481 | self.graph_state.inputs.append(graph_input) |
| 482 | |
| 483 | def handle_output(self, node: torch.fx.Node): |
| 484 | assert node.op == "output" |
| 485 | assert len(node.args) == 1, "FX.Node's args should have one arg" |
| 486 | node_args = node.args[0] |
| 487 | if isinstance(node_args, torch.fx.Node): |
| 488 | # For singleton tensor returns |
| 489 | self.graph_state.is_single_tensor_return = True |
| 490 | self.graph_state.outputs = [self.serialize_input(node_args)] |
| 491 | else: |
| 492 | assert isinstance(node_args, (tuple, list)) |