已合并
test(fx): add test for lazy_format_graph_code #35080
Lane0218创建于 5月8日
test(fx): add test for lazy_format_graph_code #35080
已合并
共 1 个文件变更+31-0
| @@ -0,0 +1,31 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch.fx lazy graph formatting APIs on NPU: | ||
| 3 | + | ||
| 4 | +1. test/test_fx.py from PyTorch community contains broad FX validations and | ||
| 5 | + triggers historical lint issues when modified, so this file is added. | ||
| 6 | +2. This file validates torch.fx._utils.lazy_format_graph_code output formatting | ||
| 7 | + for traced GraphModule instances. | ||
| 8 | +""" | ||
| 9 | + | ||
| 10 | +import torch | ||
| 11 | +from torch.fx import symbolic_trace | ||
| 12 | +from torch.fx._utils import lazy_format_graph_code | ||
| 13 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +class TestLazyFormatGraphCode(TestCase): | ||
| 17 | + def test_lazy_format_graph_code(self): | ||
| 18 | + class MyModule(torch.nn.Module): | ||
| 19 | + def forward(self, x): | ||
| 20 | + return x + 1 | ||
| 21 | + | ||
| 22 | + gm = symbolic_trace(MyModule()) | ||
| 23 | + graph_code = str(lazy_format_graph_code("fx lazy graph", gm, maybe_id=1)) | ||
| 24 | + | ||
| 25 | + self.assertIn("TRACED GRAPH", graph_code) | ||
| 26 | + self.assertIn("===== fx lazy graph 1 =====", graph_code) | ||
| 27 | + self.assertIn("def forward", graph_code) | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +if __name__ == "__main__": | ||
| 31 | + run_tests() | ||