已合并
[test] Add fx graph internal API verify tests #35606
lihaokun-2026创建于 5月14日
[test] Add fx graph internal API verify tests #35606
已合并
共 1 个文件变更+83-0
| @@ -0,0 +1,83 @@ | |||
| 1 | +# Owner(s): ["module: fx"] | ||
| 2 | +""" | ||
| 3 | +Add validation cases for selected torch.fx.graph internal APIs on NPU. | ||
| 4 | + | ||
| 5 | +Current covered APIs: | ||
| 6 | +- torch.fx.graph._format_target | ||
| 7 | +- torch.fx.graph._is_from_torch | ||
| 8 | +- torch.fx.graph._origin_type_map.get | ||
| 9 | +- torch.fx.graph._register_custom_builtin | ||
| 10 | + | ||
| 11 | +Note: | ||
| 12 | +CodeGen._gen_python_code internally defines a local helper named _format_args. | ||
| 13 | +It is not a torch.fx.graph module-level callable, so it is not added as a | ||
| 14 | +separate test case in this file. | ||
| 15 | +""" | ||
| 16 | + | ||
| 17 | +import torch_npu # noqa: F401 | ||
| 18 | + | ||
| 19 | +import torch | ||
O | |||
| 20 | +import torch.fx.graph as fx_graph | ||
| 21 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +def _test_fx_graph_custom_builtin_for_npu(x): | ||
| 25 | + return x | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +class TestFxGraphInternal(TestCase): | ||
| 29 | + def _remove_custom_builtin(self, name): | ||
| 30 | + fx_graph._custom_builtins.pop(name, None) | ||
| 31 | + fx_graph._illegal_names.pop(name, None) | ||
| 32 | + | ||
| 33 | + def test_format_target(self): | ||
| 34 | + self.assertEqual( | ||
| 35 | + fx_graph._format_target("root", "foo.bar"), | ||
| 36 | + "root.foo.bar", | ||
| 37 | + ) | ||
| 38 | + self.assertEqual( | ||
| 39 | + fx_graph._format_target("root", "foo.0"), | ||
| 40 | + 'getattr(root.foo, "0")', | ||
| 41 | + ) | ||
| 42 | + | ||
| 43 | + def test_is_from_torch(self): | ||
| 44 | + self.assertTrue(fx_graph._is_from_torch(torch.add)) | ||
| 45 | + self.assertTrue(fx_graph._is_from_torch(torch.relu)) | ||
| 46 | + | ||
| 47 | + def user_defined_func(x): | ||
| 48 | + return x | ||
| 49 | + | ||
| 50 | + self.assertFalse(fx_graph._is_from_torch(user_defined_func)) | ||
| 51 | + | ||
| 52 | + def test_origin_type_map_get(self): | ||
| 53 | + self.assertEqual(fx_graph._origin_type_map.get(list).__origin__, list) | ||
| 54 | + self.assertEqual(fx_graph._origin_type_map.get(dict).__origin__, dict) | ||
| 55 | + self.assertEqual(fx_graph._origin_type_map.get(set).__origin__, set) | ||
| 56 | + self.assertEqual(fx_graph._origin_type_map.get(tuple).__origin__, tuple) | ||
| 57 | + | ||
| 58 | + self.assertIsNone(fx_graph._origin_type_map.get(TestFxGraphInternal)) | ||
| 59 | + | ||
| 60 | + def test_register_custom_builtin(self): | ||
| 61 | + name = "_test_fx_graph_custom_builtin_for_npu" | ||
| 62 | + import_str = ( | ||
| 63 | + f"from {__name__} import _test_fx_graph_custom_builtin_for_npu" | ||
| 64 | + ) | ||
| 65 | + obj = _test_fx_graph_custom_builtin_for_npu | ||
| 66 | + | ||
| 67 | + self._remove_custom_builtin(name) | ||
| 68 | + | ||
| 69 | + fx_graph._register_custom_builtin(name, import_str, obj) | ||
| 70 | + | ||
| 71 | + self.assertIn(name, fx_graph._custom_builtins) | ||
| 72 | + self.assertEqual( | ||
| 73 | + fx_graph._custom_builtins[name].import_str, | ||
| 74 | + import_str, | ||
| 75 | + ) | ||
| 76 | + self.assertIs(fx_graph._custom_builtins[name].obj, obj) | ||
| 77 | + self.assertIs(fx_graph._illegal_names[name], obj) | ||
| 78 | + | ||
| 79 | + self._remove_custom_builtin(name) | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +if __name__ == "__main__": | ||
| 83 | + run_tests() | ||
此条代码评论区间+15至+19
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。