(self, node: torch.fx.Node)
| 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)) |
| 493 | self.graph_state.outputs = [self.serialize_input(arg) for arg in node_args] |
| 494 | |
| 495 | def serialize_operator(self, target) -> str: |
| 496 | if isinstance(target, str): |
nothing calls this directly
no test coverage detected