已合并
test(fx): add tests for graph codegen api #34462
nwww创建于 4月26日
test(fx): add tests for graph codegen api #34462
已合并
共 1 个文件变更+48-0
| @@ -0,0 +1,48 @@ | |||
| 1 | +# Owner(s): ["module: fx"] | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | +from torch.fx import symbolic_trace | ||
| 5 | +from torch.fx.graph import CodeGen, PythonCode | ||
| 6 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestFXCodegenAPI(TestCase): | ||
| 10 | + def test_graph_python_code_returns_python_code(self): | ||
| 11 | + def fn(x, y): | ||
| 12 | + return torch.relu(x + y) | ||
| 13 | + | ||
| 14 | + gm = symbolic_trace(fn) | ||
| 15 | + python_code = gm.graph.python_code("self") | ||
| 16 | + | ||
| 17 | + self.assertIsInstance(python_code, PythonCode) | ||
| 18 | + self.assertTrue(hasattr(python_code, "src")) | ||
| 19 | + self.assertTrue(hasattr(python_code, "globals")) | ||
| 20 | + self.assertIsInstance(python_code.src, str) | ||
| 21 | + self.assertIsInstance(python_code.globals, dict) | ||
| 22 | + self.assertIn("def forward", python_code.src) | ||
| 23 | + | ||
| 24 | + def test_graph_set_codegen(self): | ||
| 25 | + class ListCodeGen(CodeGen): | ||
| 26 | + def gen_fn_def(self, free_vars, maybe_return_annotation): | ||
| 27 | + return f"""def forward(self, args_list){maybe_return_annotation}: | ||
| 28 | + {", ".join(free_vars)} = args_list""" | ||
| 29 | + | ||
| 30 | + def process_inputs(self, *inputs): | ||
| 31 | + if len(inputs) != 1: | ||
| 32 | + raise RuntimeError("Expected exactly one input") | ||
| 33 | + return inputs[0] | ||
| 34 | + | ||
| 35 | + def fn(x, y): | ||
| 36 | + return x + y | ||
| 37 | + | ||
| 38 | + gm = symbolic_trace(fn) | ||
| 39 | + gm.graph.set_codegen(ListCodeGen()) | ||
| 40 | + gm.recompile() | ||
| 41 | + | ||
| 42 | + x = torch.randn(2, 3) | ||
| 43 | + y = torch.randn(2, 3) | ||
| 44 | + self.assertEqual(gm([x, y]), x + y) | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +if __name__ == "__main__": | ||
| 48 | + run_tests() | ||