已合并
test(fx): add create_args_for_root testcases into test_fx_tracer_create_arg.py #37610
q15546010075创建于 6月4日
test(fx): add create_args_for_root testcases into test_fx_tracer_create_arg.py #37610
已合并
共 1 个文件变更+117-4
| @@ -1,7 +1,13 @@ | |||
| 1 | """ | 1 | """ |
| 2 | Add validation cases for torch.fx.Tracer APIs on NPU: | 2 | Add validation cases for torch.fx.Tracer APIs on NPU: |
| 3 | -1. PyTorch community lacks sufficient and direct API validations for fx.Tracer.create_arg, so this file is added. | 3 | + |
| 4 | -2. This file validates torch.fx.Tracer.create_arg (extendable for other tracer APIs such as call_module, getattr, etc.). | 4 | +1. PyTorch community lacks sufficient and direct API validations for |
| 5 | + fx.Tracer.create_arg and fx.Tracer.create_args_for_root, so this file is added. | ||
| 6 | + | ||
| 7 | +2. This file validates the following APIs: | ||
| 8 | + - torch.fx.Tracer.create_arg | ||
| 9 | + - torch.fx.Tracer.create_args_for_root | ||
| 10 | + (extendable for other tracer APIs such as call_module, getattr, etc.) | ||
| 5 | """ | 11 | """ |
| 6 | 12 | ||
| 7 | import torch | 13 | import torch |
| @@ -11,11 +17,23 @@ import torch_npu | |||
| 11 | from torch.testing._internal.common_utils import run_tests, TestCase | 17 | from torch.testing._internal.common_utils import run_tests, TestCase |
| 12 | 18 | ||
| 13 | 19 | ||
| 20 | +class TwoArgModule(nn.Module): | ||
| 21 | + def forward(self, x, y): | ||
| 22 | + return x + y | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +class DefaultArgModule(nn.Module): | ||
| 26 | + def forward(self, x, scale=2.0): | ||
| 27 | + return x * scale | ||
| 28 | + | ||
| 29 | + | ||
| 14 | class TestTracerCreateArg(TestCase): | 30 | class TestTracerCreateArg(TestCase): |
| 15 | """ | 31 | """ |
| 16 | - Test suite for fx.Tracer.create_arg method. | 32 | + Test suite for fx.Tracer.create_arg and fx.Tracer.create_args_for_root methods. |
| 17 | Validates that create_arg correctly processes tensors, containers, | 33 | Validates that create_arg correctly processes tensors, containers, |
| 18 | - basic types, and NPU tensors during symbolic tracing. | 34 | + basic types, and NPU tensors during symbolic tracing; |
| 35 | + and that create_args_for_root correctly introspects function/module | ||
| 36 | + signatures and creates corresponding placeholder proxy nodes. | ||
| 19 | """ | 37 | """ |
| 20 | 38 | ||
| 21 | def _get_placeholder_nodes(self, graph): | 39 | def _get_placeholder_nodes(self, graph): |
| @@ -294,6 +312,101 @@ class TestTracerCreateArg(TestCase): | |||
| 294 | self.assertEqual(len(result["list"]), 2) | 312 | self.assertEqual(len(result["list"]), 2) |
| 295 | self.assertIsInstance(result["list"][0], fx.Node) | 313 | self.assertIsInstance(result["list"][0], fx.Node) |
| 296 | 314 | ||
| 315 | + # =================================================================== | ||
| 316 | + # Tests for torch.fx.Tracer.create_args_for_root | ||
| 317 | + # =================================================================== | ||
| 318 | + | ||
| 319 | + def test_create_args_for_root_basic_module(self): | ||
| 320 | + """is_module=True with single-arg Module: args = [root, placeholder('x')]""" | ||
| 321 | + | ||
| 322 | + class SimpleModule(nn.Module): | ||
| 323 | + def forward(self, x): | ||
| 324 | + return x + 1 | ||
| 325 | + | ||
| 326 | + mod = SimpleModule() | ||
| 327 | + tracer = fx.Tracer() | ||
| 328 | + tracer.root = mod | ||
| 329 | + tracer.graph = fx.Graph() | ||
| 330 | + fn, args = tracer.create_args_for_root(mod.forward, is_module=True) | ||
| 331 | + self.assertEqual(len(args), 2) | ||
| 332 | + self.assertIs(args[0], mod) | ||
| 333 | + self.assertEqual(args[1].node.op, "placeholder") | ||
| 334 | + self.assertEqual(args[1].node.name, "x") | ||
| 335 | + | ||
| 336 | + def test_create_args_for_root_two_arg_module(self): | ||
| 337 | + """is_module=True with multi-arg Module: args = [root, ph('x'), ph('y')]""" | ||
| 338 | + mod = TwoArgModule() | ||
| 339 | + tracer = fx.Tracer() | ||
| 340 | + tracer.root = mod | ||
| 341 | + tracer.graph = fx.Graph() | ||
| 342 | + fn, args = tracer.create_args_for_root(mod.forward, is_module=True) | ||
| 343 | + self.assertEqual(len(args), 3) | ||
| 344 | + self.assertIs(args[0], mod) | ||
| 345 | + self.assertEqual(args[1].node.name, "x") | ||
| 346 | + self.assertEqual(args[2].node.name, "y") | ||
| 347 | + | ||
| 348 | + def test_create_args_for_root_default_arg_module(self): | ||
| 349 | + """is_module=True with default-value param: placeholder for 'scale' created""" | ||
| 350 | + mod = DefaultArgModule() | ||
| 351 | + tracer = fx.Tracer() | ||
| 352 | + tracer.root = mod | ||
| 353 | + tracer.graph = fx.Graph() | ||
| 354 | + fn, args = tracer.create_args_for_root(mod.forward, is_module=True) | ||
| 355 | + self.assertEqual(len(args), 3) | ||
| 356 | + self.assertEqual(args[1].node.name, "x") | ||
| 357 | + self.assertEqual(args[2].node.name, "scale") | ||
| 358 | + | ||
| 359 | + def test_create_args_for_root_concrete_args_dict(self): | ||
| 360 | + """concrete_args as dict: specialise 'y' to a concrete tensor value""" | ||
| 361 | + mod = TwoArgModule() | ||
| 362 | + tracer = fx.Tracer() | ||
| 363 | + tracer.root = mod | ||
| 364 | + tracer.graph = fx.Graph() | ||
| 365 | + concrete = {"y": torch.tensor(3.0)} | ||
| 366 | + fn, args = tracer.create_args_for_root( | ||
| 367 | + mod.forward, is_module=True, concrete_args=concrete | ||
| 368 | + ) | ||
| 369 | + self.assertEqual(len(args), 3) | ||
| 370 | + self.assertIs(args[0], mod) | ||
| 371 | + | ||
| 372 | + def test_create_args_for_root_concrete_args_tuple(self): | ||
| 373 | + """concrete_args as tuple with PH: non-specialised params keep placeholder""" | ||
| 374 | + mod = TwoArgModule() | ||
| 375 | + tracer = fx.Tracer() | ||
| 376 | + tracer.root = mod | ||
| 377 | + tracer.graph = fx.Graph() | ||
| 378 | + concrete = (fx.PH, torch.tensor(5.0)) | ||
| 379 | + fn, args = tracer.create_args_for_root( | ||
| 380 | + mod.forward, is_module=True, concrete_args=concrete | ||
| 381 | + ) | ||
| 382 | + self.assertEqual(len(args), 3) | ||
| 383 | + | ||
| 384 | + def test_create_args_for_root_plain_function(self): | ||
| 385 | + """is_module=False: all params become placeholders, no self skip""" | ||
| 386 | + def my_fn(a, b): | ||
| 387 | + return a + b | ||
| 388 | + | ||
| 389 | + tracer = fx.Tracer() | ||
| 390 | + tracer.root = None | ||
| 391 | + tracer.graph = fx.Graph() | ||
| 392 | + fn, args = tracer.create_args_for_root(my_fn, is_module=False) | ||
| 393 | + self.assertEqual(len(args), 2) | ||
| 394 | + self.assertEqual(args[0].node.name, "a") | ||
| 395 | + self.assertEqual(args[1].node.name, "b") | ||
| 396 | + | ||
| 397 | + def test_create_args_for_root_function_with_defaults(self): | ||
| 398 | + """is_module=False with default-value param: placeholder created for 'b'""" | ||
| 399 | + def my_fn(a, b=10): | ||
| 400 | + return a * b | ||
| 401 | + | ||
| 402 | + tracer = fx.Tracer() | ||
| 403 | + tracer.root = None | ||
| 404 | + tracer.graph = fx.Graph() | ||
| 405 | + fn, args = tracer.create_args_for_root(my_fn, is_module=False) | ||
| 406 | + self.assertEqual(len(args), 2) | ||
| 407 | + self.assertEqual(args[0].node.name, "a") | ||
| 408 | + self.assertEqual(args[1].node.name, "b") | ||
| 409 | + | ||
| 297 | 410 | ||
| 298 | if __name__ == "__main__": | 411 | if __name__ == "__main__": |
| 299 | run_tests() | 412 | run_tests() |