(self, graph_module: torch.fx.GraphModule)
| 1334 | return arg_list |
| 1335 | |
| 1336 | def serialize_graph(self, graph_module: torch.fx.GraphModule) -> Graph: |
| 1337 | assert isinstance(graph_module, torch.fx.GraphModule) |
| 1338 | for node in graph_module.graph.nodes: |
| 1339 | try: |
| 1340 | getattr(self, f"handle_{node.op}")(node) |
| 1341 | except Exception as e: |
| 1342 | raise SerializeError( |
| 1343 | f"Failed serializing node {node} in graph: {node.format_node()}" |
| 1344 | ) from e |
| 1345 | |
| 1346 | return Graph( |
| 1347 | inputs=self.graph_state.inputs, |
| 1348 | nodes=self.graph_state.nodes, |
| 1349 | tensor_values=self.graph_state.tensor_values, |
| 1350 | sym_int_values=self.graph_state.sym_int_values, |
| 1351 | sym_bool_values=self.graph_state.sym_bool_values, |
| 1352 | custom_obj_values=self.graph_state.custom_obj_values, |
| 1353 | outputs=self.graph_state.outputs, |
| 1354 | is_single_tensor_return=self.graph_state.is_single_tensor_return, |
| 1355 | ) |
| 1356 | |
| 1357 | def serialize(self, graph_module: torch.fx.GraphModule) -> GraphModule: |
| 1358 | graph = self.serialize_graph(graph_module) |
no test coverage detected