(gm: torch.fx.GraphModule)
| 360 | |
| 361 | # TODO(zhxchen17) Try ExportPass |
| 362 | def _fixup_output_node(gm: torch.fx.GraphModule) -> None: |
| 363 | for node in reversed(gm.graph.nodes): |
| 364 | if node.op == "output": |
| 365 | with gm.graph.inserting_before(node): |
| 366 | assert len(node.args) == 1 |
| 367 | outputs = node.args[0] |
| 368 | if isinstance(outputs, torch.fx.Node): |
| 369 | val = outputs.meta.get("val") |
| 370 | if isinstance(val, list): |
| 371 | # If a list is returned, in some cases it is represented as a |
| 372 | # singular node, like `split_copy_tensor` but EXIR will return a |
| 373 | # opened-up list like `[getitem1, getitem2]` |
| 374 | outputs = [ |
| 375 | torch.fx.Proxy(outputs)[i].node for i in range(len(val)) |
| 376 | ] |
| 377 | returns, out_spec = pytree.tree_flatten(outputs) |
| 378 | node.args = (returns,) |
| 379 | return |
| 380 | |
| 381 | |
| 382 | def arrange_graph_placeholders( |
no test coverage detected