Modifies the output nodes in the submodules to return the result as a flattened list. This keeps it consistent with the result of EXIR's tracer
(gm: torch.fx.GraphModule)
| 604 | |
| 605 | |
| 606 | def flatten_output(gm: torch.fx.GraphModule) -> None: |
| 607 | """ |
| 608 | Modifies the output nodes in the submodules to return the result |
| 609 | as a flattened list. This keeps it consistent with the result of |
| 610 | EXIR's tracer |
| 611 | """ |
| 612 | for node in reversed(gm.graph.nodes): |
| 613 | if node.op == "output": |
| 614 | assert len(node.args) == 1 |
| 615 | outputs = node.args[0] |
| 616 | returns, _ = pytree.tree_flatten(outputs) |
| 617 | node.args = (returns,) |
| 618 | return |
| 619 | raise RuntimeError(f"Could not find an output node in {gm.graph}") |
| 620 | |
| 621 | |
| 622 | def _default_decomposition_table( |