| 24 | # this transform assumes that all operations take in only one input and return |
| 25 | # one output. |
| 26 | def invert(model: torch.nn.Module) -> torch.nn.Module: |
| 27 | fx_model = fx.symbolic_trace(model) |
| 28 | new_graph = fx.Graph() # As we're building up a new graph |
| 29 | env = {} |
| 30 | for node in reversed(fx_model.graph.nodes): |
| 31 | if node.op == 'call_function': |
| 32 | # This creates a node in the new graph with the inverse function, |
| 33 | # and passes `env[node.name]` (i.e. the previous output node) as |
| 34 | # input. |
| 35 | new_node = new_graph.call_function(invert_mapping[node.target], (env[node.name],)) |
| 36 | env[node.args[0].name] = new_node |
| 37 | elif node.op == 'output': |
| 38 | # We turn the output into an input placeholder |
| 39 | new_node = new_graph.placeholder(node.name) |
| 40 | env[node.args[0].name] = new_node |
| 41 | elif node.op == 'placeholder': |
| 42 | # We turn the input placeholder into an output |
| 43 | new_graph.output(env[node.name]) |
| 44 | else: |
| 45 | raise RuntimeError("Not implemented") |
| 46 | |
| 47 | new_graph.lint() |
| 48 | return fx.GraphModule(fx_model, new_graph) |
| 49 | |
| 50 | |
| 51 | def f(x): |