(self, graph_module: torch.fx.GraphModule)
| 1014 | |
| 1015 | class TestPass(ExportPass): |
| 1016 | def call(self, graph_module: torch.fx.GraphModule) -> PassResult: |
| 1017 | permute_dims = [1, 0, 2] |
| 1018 | for node in graph_module.graph.nodes: |
| 1019 | if node.op == "placeholder" and str(node) == "a": |
| 1020 | inverse_dims = [ |
| 1021 | permute_dims.index(x) for x in range(len(permute_dims)) |
| 1022 | ] |
| 1023 | |
| 1024 | with graph_module.graph.inserting_after(node): |
| 1025 | permute = graph_module.graph.call_function( |
| 1026 | exir_ops.edge.aten.permute_copy.default, |
| 1027 | args=(node, inverse_dims), |
| 1028 | ) |
| 1029 | permute.meta = node.meta.copy() |
| 1030 | node.meta["val"] = node.meta["val"].permute(permute_dims) |
| 1031 | node.replace_all_uses_with( |
| 1032 | permute, lambda x, permute=permute: x is not permute |
| 1033 | ) |
| 1034 | break |
| 1035 | return PassResult(graph_module, True) |
| 1036 | |
| 1037 | edge = edge.transform([TestPass()]) |
| 1038 | et = edge.to_executorch() |
nothing calls this directly
no test coverage detected