| 68 | # during graph manipulation. |
| 69 | |
| 70 | def inline_lowp_func(n : torch.fx.Node): |
| 71 | # If we find a call to a function in our "lowp" module, inline it |
| 72 | if n.op == 'call_function' and n.target.__module__ == inline_lowp_func.__module__: |
| 73 | # We want to insert the operations comprising the implementation of the |
| 74 | # function before the function itself. Then, we can swap the output value |
| 75 | # of the function call with the output value for its implementation nodes |
| 76 | tracer = torch.fx.proxy.GraphAppendingTracer(n.graph) |
| 77 | with n.graph.inserting_before(n): |
| 78 | # We can inline code by using `fx.Proxy` instances. |
| 79 | # map_arg traverses all aggregate types and applies the given function |
| 80 | # to Node instances in the data structure. In this case, we are applying |
| 81 | # the fx.Proxy constructor. |
| 82 | proxy_args = torch.fx.node.map_arg(n.args, lambda x: torch.fx.Proxy(x, tracer)) |
| 83 | proxy_kwargs = torch.fx.node.map_arg(n.kwargs, lambda x: torch.fx.Proxy(x, tracer)) |
| 84 | # Call the function itself with proxy arguments. This will emit |
| 85 | # nodes in the graph corresponding to the operations in the im- |
| 86 | # plementation of the function |
| 87 | output_proxy = n.target(*proxy_args, **proxy_kwargs) |
| 88 | # Now replace the original node's uses with the output node of |
| 89 | # the implementation. |
| 90 | node.replace_all_uses_with(output_proxy.node) |
| 91 | # Delete the old node |
| 92 | node.graph.erase_node(node) |
| 93 | |
| 94 | for node in traced.graph.nodes: |
| 95 | if node.op == 'call_function' and node.target is sigmoid_lowp: |