已合并
Synchronize official version 2.2 code for test cases. #8894
xiaxia3创建于 2024年1月12日
Synchronize official version 2.2 code for test cases. #8894
已合并
xiaxia3创建于 2024年1月12日
refs/pull/8894/head合入到master
9 个文件变更+357-117
@@ -1901,6 +1901,7 @@ class TestConvolutionNNDeviceType(NNTestCase):
1901 1901 
1902 @onlyPRIVATEUSE11902 @onlyPRIVATEUSE1
1903 @largeTensorTest('12GB')1903 @largeTensorTest('12GB')
1904+ @skipIfRocmVersionLessThan((6, 0))
1904 def test_conv_transposed_large(self, device):1905 def test_conv_transposed_large(self, device):
1905 dtype = torch.half if self.device_type == 'npu' else torch.float1906 dtype = torch.half if self.device_type == 'npu' else torch.float
1906 conv = nn.ConvTranspose2d(1, 1, 1, 1, bias=False).to(device).to(dtype)1907 conv = nn.ConvTranspose2d(1, 1, 1, 1, bias=False).to(device).to(dtype)
@@ -2031,6 +2032,7 @@ class TestConvolutionNNDeviceType(NNTestCase):
2031 @dtypesIfPRIVATEUSE1(*floating_types_and(torch.half, *[torch.bfloat16] if AMPERE_OR_ROCM else []))2032 @dtypesIfPRIVATEUSE1(*floating_types_and(torch.half, *[torch.bfloat16] if AMPERE_OR_ROCM else []))
2032 @dtypes(torch.float)2033 @dtypes(torch.float)
2033 @torch.backends.cudnn.flags(enabled=True, benchmark=False)2034 @torch.backends.cudnn.flags(enabled=True, benchmark=False)
2035+ @unittest.skipIf(TEST_WITH_ROCM, "Skipped on ROCm, since it is failing on ROCm 5.7")
2034 def test_Conv2d_naive_groups(self, device, dtype):2036 def test_Conv2d_naive_groups(self, device, dtype):
2035 # Check that grouped convolutions matches two half convolutions2037 # Check that grouped convolutions matches two half convolutions
2036 m = nn.Conv2d(4, 4, kernel_size=3, groups=2).to(device, dtype)2038 m = nn.Conv2d(4, 4, kernel_size=3, groups=2).to(device, dtype)
@@ -11,7 +11,9 @@ import torch.nn.init as init
11import torch.nn.functional as F11import torch.nn.functional as F
12import torch_npu12import torch_npu
13import torch_npu.testing13import torch_npu.testing
14-from torch.testing._internal.common_utils import TestCase, TEST_SCIPY, skipIfNoLapack14+from torch.testing._internal.common_utils import (
15+ TestCase, TEST_SCIPY, skipIfNoLapack, skipIfTorchDynamo, run_tests
16+)
15 17 
16if TEST_SCIPY:18if TEST_SCIPY:
17 from scipy import stats19 from scipy import stats
@@ -103,6 +105,7 @@ class TestNNInit(TestCase):
103 assert self._is_uniform(input_tensor, a, b)105 assert self._is_uniform(input_tensor, a, b)
104 106 
105 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")107 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
108+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
106 def test_normal(self):109 def test_normal(self):
107 for dims in [1, 2, 4]:110 for dims in [1, 2, 4]:
108 input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)111 input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)
@@ -113,6 +116,7 @@ class TestNNInit(TestCase):
113 assert self._is_normal(input_tensor, mean, std)116 assert self._is_normal(input_tensor, mean, std)
114 117 
115 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")118 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
119+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
116 def test_trunc_normal(self):120 def test_trunc_normal(self):
117 for dims in [1, 2, 4]:121 for dims in [1, 2, 4]:
118 input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)122 input_tensor = self._create_random_nd_tensor(dims, size_min=30, size_max=50)
@@ -125,6 +129,7 @@ class TestNNInit(TestCase):
125 assert self._is_trunc_normal(input_tensor, mean, std, a, b)129 assert self._is_trunc_normal(input_tensor, mean, std, a, b)
126 130 
127 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")131 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
132+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
128 def test_trunc_normal_generator(self):133 def test_trunc_normal_generator(self):
129 gen = torch.Generator()134 gen = torch.Generator()
130 gen.manual_seed(42)135 gen.manual_seed(42)
@@ -279,6 +284,7 @@ class TestNNInit(TestCase):
279 assert self._is_uniform(input_tensor, -bounds, bounds)284 assert self._is_uniform(input_tensor, -bounds, bounds)
280 285 
281 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")286 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
287+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
282 def test_xavier_normal(self):288 def test_xavier_normal(self):
283 for use_gain in [True, False]:289 for use_gain in [True, False]:
284 for dims in [2, 4]:290 for dims in [2, 4]:
@@ -323,6 +329,7 @@ class TestNNInit(TestCase):
323 _ = init.kaiming_normal_(tensor)329 _ = init.kaiming_normal_(tensor)
324 330 
325 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")331 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
332+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
326 def test_kaiming_uniform(self):333 def test_kaiming_uniform(self):
327 for use_a in [True, False]:334 for use_a in [True, False]:
328 for dims in [2, 4]:335 for dims in [2, 4]:
@@ -351,6 +358,7 @@ class TestNNInit(TestCase):
351 assert self._is_uniform(input_tensor, -bounds, bounds)358 assert self._is_uniform(input_tensor, -bounds, bounds)
352 359 
353 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")360 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
361+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
354 def test_kaiming_normal(self):362 def test_kaiming_normal(self):
355 for use_a in [True, False]:363 for use_a in [True, False]:
356 for dims in [2, 4]:364 for dims in [2, 4]:
@@ -385,6 +393,7 @@ class TestNNInit(TestCase):
385 init.sparse_(tensor, sparsity)393 init.sparse_(tensor, sparsity)
386 394 
387 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")395 @unittest.skipIf(not TEST_SCIPY, "Scipy not found.")
396+ @skipIfTorchDynamo("scipy.kstest is failing under dynamo")
388 def test_sparse_default_std(self):397 def test_sparse_default_std(self):
389 for use_random_std in [True, False]:398 for use_random_std in [True, False]:
390 input_tensor = self._create_random_nd_tensor(2, size_min=30, size_max=35)399 input_tensor = self._create_random_nd_tensor(2, size_min=30, size_max=35)
@@ -434,3 +443,7 @@ class TestNNInit(TestCase):
434 443 
435 with self.assertWarnsRegex(UserWarning, 'deprecated', msg='methods not suffixed with underscore should be deprecated'):444 with self.assertWarnsRegex(UserWarning, 'deprecated', msg='methods not suffixed with underscore should be deprecated'):
436 fn()445 fn()
446+ 
447+ 
448+if __name__ == '__main__':
449+ run_tests()
@@ -189,7 +189,6 @@ class DummyContextManager:
189 189 
190 190 
191class TestModuleHooks(TestCase):191class TestModuleHooks(TestCase):
192- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
193 @parametrize_test("named_tuple", (True, False))192 @parametrize_test("named_tuple", (True, False))
194 def test_forward_hooks(self, named_tuple):193 def test_forward_hooks(self, named_tuple):
195 fired_hooks: List[int] = []194 fired_hooks: List[int] = []
@@ -212,7 +211,6 @@ class TestModuleHooks(TestCase):
212 model(x)[0].sum().backward()211 model(x)[0].sum().backward()
213 self.assertEqual(fired_hooks, expected + expected)212 self.assertEqual(fired_hooks, expected + expected)
214 213 
215- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
216 @parametrize_test("named_tuple", (True, False))214 @parametrize_test("named_tuple", (True, False))
217 def test_forward_pre_hooks(self, named_tuple):215 def test_forward_pre_hooks(self, named_tuple):
218 fired_hooks: List[int] = []216 fired_hooks: List[int] = []
@@ -239,7 +237,6 @@ class TestModuleHooks(TestCase):
239 model(x)[0].sum().backward()237 model(x)[0].sum().backward()
240 self.assertEqual(fired_hooks, expected + expected)238 self.assertEqual(fired_hooks, expected + expected)
241 239 
242- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
243 @parametrize_test("named_tuple", (True, False))240 @parametrize_test("named_tuple", (True, False))
244 def test_full_backward_hooks(self, named_tuple):241 def test_full_backward_hooks(self, named_tuple):
245 fired_hooks: List[int] = []242 fired_hooks: List[int] = []
@@ -262,7 +259,6 @@ class TestModuleHooks(TestCase):
262 model(x)[0].sum().backward()259 model(x)[0].sum().backward()
263 self.assertEqual(fired_hooks, expected + expected)260 self.assertEqual(fired_hooks, expected + expected)
264 261 
265- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
266 @parametrize_test("named_tuple", (True, False))262 @parametrize_test("named_tuple", (True, False))
267 def test_full_backward_pre_hooks(self, named_tuple):263 def test_full_backward_pre_hooks(self, named_tuple):
268 fired_hooks: List[int] = []264 fired_hooks: List[int] = []
@@ -290,19 +286,23 @@ class TestModuleHooks(TestCase):
290 self.assertEqual(fired_hooks, expected + expected)286 self.assertEqual(fired_hooks, expected + expected)
291 287 
292 # Backward pre hook can affect subsequent gradient computation288 # Backward pre hook can affect subsequent gradient computation
293- a = torch.ones(2, requires_grad=True)289+ for rg in [True, False]:
294- model = nn.Linear(2, 2)290+ a = torch.ones(2, requires_grad=rg)
291+ model = nn.Linear(2, 2)
295 292 
296- def fn(_unused_module, grad_output):293+ def fn(_unused_module, grad_output):
297- return (grad_output[0] * 0,)294+ return (grad_output[0] * 0,)
298 295 
299- model.register_full_backward_pre_hook(fn)296+ model.register_full_backward_pre_hook(fn)
300 297 
301- out = model(a)298+ out = model(a)
302- out.sum().backward()299+ out.sum().backward()
303- self.assertEqual(a.grad, torch.zeros_like(a))300+ self.assertEqual(model.weight.grad, torch.zeros(2, 2))
301+ if rg:
302+ self.assertEqual(a.grad, torch.zeros_like(a))
303+ else:
304+ self.assertIsNone(a.grad)
304 305 
305- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
306 @parametrize_test("named_tuple", (True, False))306 @parametrize_test("named_tuple", (True, False))
307 def test_mixed_hooks(self, named_tuple):307 def test_mixed_hooks(self, named_tuple):
308 fired_hooks: List[int] = []308 fired_hooks: List[int] = []
@@ -330,7 +330,6 @@ class TestModuleHooks(TestCase):
330 model(x)[0].sum().backward()330 model(x)[0].sum().backward()
331 self.assertEqual(fired_hooks, [0, 1, 2, 3, 0, 1, 2, 3])331 self.assertEqual(fired_hooks, [0, 1, 2, 3, 0, 1, 2, 3])
332 332 
333- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
334 def test_kwarg_hooks(self):333 def test_kwarg_hooks(self):
335 # 1. test forward pre hook334 # 1. test forward pre hook
336 fired_hooks: List[int] = []335 fired_hooks: List[int] = []
@@ -386,7 +385,6 @@ class TestModuleHooks(TestCase):
386 out = model(x, bias=bias)385 out = model(x, bias=bias)
387 self.assertEqual(out, x + 2 * bias, rtol=0, atol=1e-5)386 self.assertEqual(out, x + 2 * bias, rtol=0, atol=1e-5)
388 387 
389- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
390 def test_remove_kwarg_hooks(self):388 def test_remove_kwarg_hooks(self):
391 # test forward pre and forward hooks389 # test forward pre and forward hooks
392 fired_hooks: List[int] = []390 fired_hooks: List[int] = []
@@ -432,7 +430,6 @@ class TestModuleHooks(TestCase):
432 forward_pre_hook_handle.id in model._forward_pre_hooks_with_kwargs430 forward_pre_hook_handle.id in model._forward_pre_hooks_with_kwargs
433 )431 )
434 432 
435- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
436 def test_always_called_forward_hooks(self):433 def test_always_called_forward_hooks(self):
437 x: torch.Tensor = torch.ones(10, 10)434 x: torch.Tensor = torch.ones(10, 10)
438 model = FailsInForwardModel()435 model = FailsInForwardModel()
@@ -520,7 +517,6 @@ class TestModuleHooks(TestCase):
520 model(x)517 model(x)
521 self.assertEqual(stack, [2, -1, 2, -1, 2, -1, 2, -1, 2, -1, 2, -1, 2])518 self.assertEqual(stack, [2, -1, 2, -1, 2, -1, 2, -1, 2, -1, 2, -1, 2])
522 519 
523- @skipIfTorchDynamo("Dynamo does not yet capture hooks")
524 def test_bw_hook_warning_for_non_tensor_or_tuple(self):520 def test_bw_hook_warning_for_non_tensor_or_tuple(self):
525 # Test to verify that backward hook raises warning521 # Test to verify that backward hook raises warning
526 # if result is not a Tensor or tuple of Tensors.522 # if result is not a Tensor or tuple of Tensors.
@@ -890,7 +886,6 @@ class TestModuleGlobalHooks(TestCase):
890 expected_grad = sig_x * (1 - sig_x) * 2886 expected_grad = sig_x * (1 - sig_x) * 2
891 self.assertEqual(input1.grad, expected_grad)887 self.assertEqual(input1.grad, expected_grad)
892 888 
893- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
894 def test_module_global_forward_preforward_hook_writeable(self):889 def test_module_global_forward_preforward_hook_writeable(self):
895 module = nn.Sigmoid()890 module = nn.Sigmoid()
896 input1 = torch.randn(5, 5, requires_grad=True)891 input1 = torch.randn(5, 5, requires_grad=True)
@@ -912,7 +907,6 @@ class TestModuleGlobalHooks(TestCase):
912 expected_grad = -sig_x * (1 - sig_x) * 2 * mask907 expected_grad = -sig_x * (1 - sig_x) * 2 * mask
913 self.assertEqual(input1.grad, expected_grad)908 self.assertEqual(input1.grad, expected_grad)
914 909 
915- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
916 def test_module_forward_preforward_hook_removable(self):910 def test_module_forward_preforward_hook_removable(self):
917 """911 """
918 This test is to test when multiple pre-forward hook functions can be912 This test is to test when multiple pre-forward hook functions can be
@@ -948,7 +942,6 @@ class TestModuleGlobalHooks(TestCase):
948 self.assertEqual(len(handle.hooks_dict_ref()), 0)942 self.assertEqual(len(handle.hooks_dict_ref()), 0)
949 self.assertEqual(len(handle_2.hooks_dict_ref()), 0)943 self.assertEqual(len(handle_2.hooks_dict_ref()), 0)
950 944 
951- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
952 def test_module_forward_forward_hook_removable(self):945 def test_module_forward_forward_hook_removable(self):
953 """946 """
954 This test is to test when multiple forward hook functions can be registered947 This test is to test when multiple forward hook functions can be registered
@@ -1133,7 +1126,6 @@ class TestModuleHookNN(NNTestCase):
1133 test_fwd.remove()1126 test_fwd.remove()
1134 test_bwd.remove()1127 test_bwd.remove()
1135 1128 
1136- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
1137 def test_hooks(self):1129 def test_hooks(self):
1138 self._test_hooks("register_backward_hook")1130 self._test_hooks("register_backward_hook")
1139 self._test_hooks("register_full_backward_hook")1131 self._test_hooks("register_full_backward_hook")
@@ -1151,7 +1143,6 @@ class TestModuleHookNN(NNTestCase):
1151 output = bn(torch.randn(5, 5, requires_grad=True))1143 output = bn(torch.randn(5, 5, requires_grad=True))
1152 output.sum().backward()1144 output.sum().backward()
1153 1145 
1154- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
1155 def test_backward_hooks_interaction(self):1146 def test_backward_hooks_interaction(self):
1156 # Test to make sure that the grad_outputs1147 # Test to make sure that the grad_outputs
1157 # updated by full_backward_pre_hook are received by1148 # updated by full_backward_pre_hook are received by
@@ -1226,7 +1217,6 @@ class TestModuleHookNN(NNTestCase):
1226 mod.register_full_backward_hook(lambda mod, gI, gO: None)1217 mod.register_full_backward_hook(lambda mod, gI, gO: None)
1227 mod(inp, inp.detach(), inp)1218 mod(inp, inp.detach(), inp)
1228 1219 
1229- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
1230 def test_hook_no_requires_grad(self):1220 def test_hook_no_requires_grad(self):
1231 mod = nn.Linear(2, 3)1221 mod = nn.Linear(2, 3)
1232 1222 
@@ -1420,7 +1410,6 @@ class TestModuleHookNN(NNTestCase):
1420 with module.register_full_backward_hook(bw_hook):1410 with module.register_full_backward_hook(bw_hook):
1421 module(inp1, inp2).sum().backward()1411 module(inp1, inp2).sum().backward()
1422 1412 
1423- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
1424 def test_hook_backward_writeable(self):1413 def test_hook_backward_writeable(self):
1425 module = nn.Sigmoid()1414 module = nn.Sigmoid()
1426 input1 = torch.randn(5, 5, requires_grad=True)1415 input1 = torch.randn(5, 5, requires_grad=True)
@@ -1438,7 +1427,6 @@ class TestModuleHookNN(NNTestCase):
1438 expected_grad = sig_x * (1 - sig_x) * 21427 expected_grad = sig_x * (1 - sig_x) * 2
1439 self.assertEqual(input1.grad, expected_grad)1428 self.assertEqual(input1.grad, expected_grad)
1440 1429 
1441- @skipIfTorchDynamo("TorchDynamo does not work well with hooks")
1442 def test_hook_forward_preforward_writable(self):1430 def test_hook_forward_preforward_writable(self):
1443 module = nn.Sigmoid()1431 module = nn.Sigmoid()
1444 input1 = torch.randn(5, 5, requires_grad=True)1432 input1 = torch.randn(5, 5, requires_grad=True)
@@ -13,6 +13,7 @@ import torch_npu
13import torch_npu.testing13import torch_npu.testing
14from torch.testing._internal.common_utils import run_tests, skipIfNoLapack, \14from torch.testing._internal.common_utils import run_tests, skipIfNoLapack, \
15 TemporaryFileName, instantiate_parametrized_tests, set_default_dtype15 TemporaryFileName, instantiate_parametrized_tests, set_default_dtype
16+from torch.testing._internal.common_device_type import instantiate_device_type_tests
16from torch.testing._internal.common_nn import NNTestCase17from torch.testing._internal.common_nn import NNTestCase
17from torch.testing._internal.common_utils import gradcheck, TEST_PRIVATEUSE118from torch.testing._internal.common_utils import gradcheck, TEST_PRIVATEUSE1
18 19 
@@ -1505,35 +1506,6 @@ class TestNNParametrization(NNTestCase):
1505 m.weight = torch.randn(5, 5)1506 m.weight = torch.randn(5, 5)
1506 torch.nn.utils.parametrize.remove_parametrizations(m, "weight")1507 torch.nn.utils.parametrize.remove_parametrizations(m, "weight")
1507 1508 
1508- def test_weight_norm_parametrization(self):
1509- for dtype in [torch.float, torch.bfloat16]:
1510- input1 = torch.randn(3, 4, dtype=dtype)
1511- m = nn.Linear(4, 5).to(dtype=dtype)
1512- expected_output = m(input1)
1513- 
1514- # add weight normalization
1515- m = torch.nn.utils.parametrizations.weight_norm(m)
1516- self.assertEqual(m.parametrizations.weight.original1.size(), m.weight.size())
1517- self.assertEqual(m.parametrizations.weight.original0.size(), (5, 1))
1518- self.assertEqual(m(input1), expected_output)
1519- 
1520- # remove weight norm
1521- torch.nn.utils.parametrize.remove_parametrizations(m, "weight")
1522- self.assertFalse(hasattr(m, "parametrizations"))
1523- self.assertEqual(m(input1), expected_output)
1524- 
1525- # test with dim=1
1526- m = torch.nn.utils.parametrizations.weight_norm(m, dim=1)
1527- self.assertEqual(m.parametrizations.weight.original1.size(), m.weight.size())
1528- self.assertEqual(m.parametrizations.weight.original0.size(), (1, 4))
1529- self.assertEqual(m(input1), expected_output)
1530- 
1531- # test with dim=None
1532- m = nn.Linear(4, 5).to(dtype=dtype)
1533- expected_output = m(input1)
1534- m = torch.nn.utils.parametrizations.weight_norm(m, dim=None)
1535- self.assertEqual(m(input1), expected_output)
1536- 
1537 def test_weight_norm_state_dict_compat(self):1509 def test_weight_norm_state_dict_compat(self):
1538 m = nn.Linear(4, 5)1510 m = nn.Linear(4, 5)
1539 m = torch.nn.utils.weight_norm(m)1511 m = torch.nn.utils.weight_norm(m)
@@ -1560,6 +1532,39 @@ class TestNNParametrization(NNTestCase):
1560 self.assertEqual(m(input1), m2(input1))1532 self.assertEqual(m(input1), m2(input1))
1561 1533 
1562 1534 
1535+class TestNNParametrizationDevice(NNTestCase):
1536+ def test_weight_norm_parametrization(self, device):
1537+ for dtype in [torch.float, torch.bfloat16]:
1538+ input1 = torch.randn(3, 4, dtype=dtype, device=device)
1539+ m = nn.Linear(4, 5, dtype=dtype, device=device)
1540+ expected_output = m(input1)
1541+ 
1542+ # add weight normalization
1543+ m = torch.nn.utils.parametrizations.weight_norm(m)
1544+ self.assertEqual(m.parametrizations.weight.original1.size(), m.weight.size())
1545+ self.assertEqual(m.parametrizations.weight.original0.size(), (5, 1))
1546+ self.assertEqual(m(input1), expected_output)
1547+ 
1548+ # remove weight norm
1549+ torch.nn.utils.parametrize.remove_parametrizations(m, "weight")
1550+ self.assertFalse(hasattr(m, "parametrizations"))
1551+ self.assertEqual(m(input1), expected_output)
1552+ 
1553+ # test with dim=1
1554+ m = torch.nn.utils.parametrizations.weight_norm(m, dim=1)
1555+ self.assertEqual(m.parametrizations.weight.original1.size(), m.weight.size())
1556+ self.assertEqual(m.parametrizations.weight.original0.size(), (1, 4))
1557+ self.assertEqual(m(input1), expected_output)
1558+ 
1559+ # test with dim=None
1560+ m = nn.Linear(4, 5, dtype=dtype, device=device)
1561+ expected_output = m(input1)
1562+ m = torch.nn.utils.parametrizations.weight_norm(m, dim=None)
1563+ self.assertEqual(m(input1), expected_output)
1564+ 
1565+ 
1566+only_for = ("cpu", "npu")
1567+instantiate_device_type_tests(TestNNParametrizationDevice, globals(), only_for=only_for)
1563instantiate_parametrized_tests(TestNNParametrization)1568instantiate_parametrized_tests(TestNNParametrization)
1564 1569 
1565if __name__ == '__main__':1570if __name__ == '__main__':
@@ -9,6 +9,7 @@ import os
9import random9import random
10import itertools10import itertools
11import math11import math
12+import operator
12 13 
13from torch import inf, nan14from torch import inf, nan
14import torch15import torch
@@ -43,11 +44,11 @@ class TestAvgPool(TestCase):
43 return joined_x.view(1, joined_x.numel())44 return joined_x.view(1, joined_x.numel())
44 45 
45 def _avg_pool2d(self, x, kernel_size):46 def _avg_pool2d(self, x, kernel_size):
46- size = reduce((lambda x, y: x * y), kernel_size)47+ size = reduce(operator.mul, kernel_size)
47 return self._sum_pool2d(x, kernel_size) / size48 return self._sum_pool2d(x, kernel_size) / size
48 49 
49 def _avg_pool3d(self, x, kernel_size):50 def _avg_pool3d(self, x, kernel_size):
50- size = reduce((lambda x, y: x * y), kernel_size)51+ size = reduce(operator.mul, kernel_size)
51 return self._sum_pool3d(x, kernel_size) / size52 return self._sum_pool3d(x, kernel_size) / size
52 53 
53 def test_doubletensor_avg_pool2d(self):54 def test_doubletensor_avg_pool2d(self):
@@ -209,13 +210,13 @@ class TestPoolingNN(NNTestCase):
209 self.assertEqual(out, ref_out)210 self.assertEqual(out, ref_out)
210 self.assertEqual(input1.grad, ref_input.grad)211 self.assertEqual(input1.grad, ref_input.grad)
211 212 
212- def test_adaptive_pooling_bfloat16(self):213+ def test_adaptive_pooling_lower_precision(self):
213- def _test_adaptive_pooling_bfloat16(self, device, mod, memory_format):214+ def _test_adaptive_pooling_lower_precision(self, device, dtype, mod, memory_format):
214 input1 = torch.randint(1, 10, (3, 19, 8, 8), dtype=torch.float32)215 input1 = torch.randint(1, 10, (3, 19, 8, 8), dtype=torch.float32)
215 input1 = input1.to(device).to(memory_format=memory_format).requires_grad_()216 input1 = input1.to(device).to(memory_format=memory_format).requires_grad_()
216 pool = mod((7, 7)).to(device)217 pool = mod((7, 7)).to(device)
217 218 
218- input2 = input1.detach().clone().bfloat16().requires_grad_(True)219+ input2 = input1.detach().clone().to(dtype=dtype).requires_grad_(True)
219 220 
220 out = pool(input1)221 out = pool(input1)
221 out.sum().backward()222 out.sum().backward()
@@ -223,17 +224,18 @@ class TestPoolingNN(NNTestCase):
223 out2.sum().backward()224 out2.sum().backward()
224 225 
225 self.assertTrue(out2.is_contiguous(memory_format=memory_format))226 self.assertTrue(out2.is_contiguous(memory_format=memory_format))
226- self.assertEqual(out2.dtype, torch.bfloat16)227+ self.assertEqual(out2.dtype, dtype)
227- self.assertEqual(input2.grad.dtype, torch.bfloat16)228+ self.assertEqual(input2.grad.dtype, dtype)
228 self.assertEqual(out, out2.float(), atol=0.1, rtol=0)229 self.assertEqual(out, out2.float(), atol=0.1, rtol=0)
229 self.assertEqual(input1.grad, input2.grad.float(), atol=0.1, rtol=0)230 self.assertEqual(input1.grad, input2.grad.float(), atol=0.1, rtol=0)
230 231 
231 device_list = ['cpu']232 device_list = ['cpu']
232 for device in device_list:233 for device in device_list:
233- _test_adaptive_pooling_bfloat16(self, device, torch.nn.AdaptiveAvgPool2d, torch.contiguous_format)234+ for dtype in [torch.bfloat16, torch.float16]:
234- _test_adaptive_pooling_bfloat16(self, device, torch.nn.AdaptiveAvgPool2d, torch.channels_last)235+ _test_adaptive_pooling_lower_precision(self, device, dtype, torch.nn.AdaptiveAvgPool2d, torch.contiguous_format)
235- _test_adaptive_pooling_bfloat16(self, device, torch.nn.AdaptiveMaxPool2d, torch.contiguous_format)236+ _test_adaptive_pooling_lower_precision(self, device, dtype, torch.nn.AdaptiveAvgPool2d, torch.channels_last)
236- _test_adaptive_pooling_bfloat16(self, device, torch.nn.AdaptiveMaxPool2d, torch.channels_last)237+ _test_adaptive_pooling_lower_precision(self, device, dtype, torch.nn.AdaptiveMaxPool2d, torch.contiguous_format)
238+ _test_adaptive_pooling_lower_precision(self, device, dtype, torch.nn.AdaptiveMaxPool2d, torch.channels_last)
237 239 
238 @unittest.skipIf(not TEST_PRIVATEUSE1, "NPU unavailable")240 @unittest.skipIf(not TEST_PRIVATEUSE1, "NPU unavailable")
239 @largeTensorTest('12GB', device='npu')241 @largeTensorTest('12GB', device='npu')
@@ -377,6 +379,14 @@ class TestPoolingNN(NNTestCase):
377 with self.assertRaises(RuntimeError):379 with self.assertRaises(RuntimeError):
378 F.max_unpool3d(x, torch.zeros(x.shape, dtype=int), [1, 1])380 F.max_unpool3d(x, torch.zeros(x.shape, dtype=int), [1, 1])
379 381 
382+ def test_quantized_max_pool1d_empty_kernel(self):
383+ # This used to segfault when called with an empty kernel
384+ # see pytorch issues 116323
385+ base = torch.randn(1)
386+ temp_tensor = torch.quantize_per_tensor(base, 0.1, 10, torch.quint2x4)
387+ with self.assertRaises(RuntimeError):
388+ torch.quantized_max_pool1d(temp_tensor, [])
389+ 
380 390 
381class TestPoolingNNDeviceType(NNTestCase):391class TestPoolingNNDeviceType(NNTestCase):
382 @onlyNativeDeviceTypes392 @onlyNativeDeviceTypes
@@ -1031,9 +1041,10 @@ torch_npu.npu.synchronize()
1031 helper(None, 3, 50, 50, ks=5)1041 helper(None, 3, 50, 50, ks=5)
1032 1042 
1033 @onlyCPU1043 @onlyCPU
1034- def test_avg_pool2d_bfloat16(self, device):1044+ @dtypes(torch.half, torch.bfloat16)
1045+ def test_avg_pool2d_reduced_floating(self, device, dtype):
1035 def helper(n, c, h, w, kernel_size, stride, memory_format):1046 def helper(n, c, h, w, kernel_size, stride, memory_format):
1036- input1 = torch.randn(n, c, h, w, dtype=torch.float32, device=device).bfloat16()1047+ input1 = torch.randn(n, c, h, w, dtype=torch.float32, device=device).to(dtype=dtype)
1037 input1 = input1.to(memory_format=memory_format).requires_grad_()1048 input1 = input1.to(memory_format=memory_format).requires_grad_()
1038 pool = torch.nn.AvgPool2d(kernel_size, stride).to(device)1049 pool = torch.nn.AvgPool2d(kernel_size, stride).to(device)
1039 1050 
@@ -1045,10 +1056,10 @@ torch_npu.npu.synchronize()
1045 out2.sum().backward()1056 out2.sum().backward()
1046 1057 
1047 self.assertTrue(out.is_contiguous(memory_format=memory_format))1058 self.assertTrue(out.is_contiguous(memory_format=memory_format))
1048- self.assertEqual(out.dtype, torch.bfloat16)1059+ self.assertEqual(out.dtype, dtype)
1049- self.assertEqual(input1.grad.dtype, torch.bfloat16)1060+ self.assertEqual(input1.grad.dtype, dtype)
1050- self.assertEqual(out, out2.bfloat16())1061+ self.assertEqual(out, out2.to(dtype=dtype))
1051- self.assertEqual(input1.grad, input2.grad.bfloat16())1062+ self.assertEqual(input1.grad, input2.grad.to(dtype=dtype))
1052 1063 
1053 helper(4, 30, 8, 8, 7, 1, torch.contiguous_format)1064 helper(4, 30, 8, 8, 7, 1, torch.contiguous_format)
1054 helper(4, 65, 8, 8, 7, 1, torch.channels_last)1065 helper(4, 65, 8, 8, 7, 1, torch.channels_last)
@@ -1057,38 +1068,48 @@ torch_npu.npu.synchronize()
1057 1068 
1058 @dtypes(torch.float, torch.double)1069 @dtypes(torch.float, torch.double)
1059 def test_adaptive_pooling_max_nhwc(self, device, dtype):1070 def test_adaptive_pooling_max_nhwc(self, device, dtype):
1060- def helper(n, c, h, w, output_height, output_width, contig):1071+ def helper(input_size, output_plane_size, contig):
1061- input1 = torch.randint(1, 10, (n, c, h, w), device=device, dtype=dtype)1072+ n_plane_dims = len(output_plane_size)
1062- input1 = input1.contiguous(memory_format=torch.channels_last)1073+ mod = torch.nn.AdaptiveMaxPool2d if n_plane_dims == 2 else torch.nn.AdaptiveMaxPool3d
1063- grad = torch.randint(1, 10, (4, 8, output_height, output_width), device=device, dtype=dtype)1074+ channels_last = torch.channels_last if n_plane_dims == 2 else torch.channels_last_3d
1064- grad = grad.contiguous(memory_format=torch.channels_last)1075+ output_size = input_size[:2] + output_plane_size
1076+ input1 = torch.randint(1, 10, input_size, device=device, dtype=dtype)
1077+ input1 = input1.contiguous(memory_format=channels_last)
1078+ grad = torch.randint(1, 10, output_size, device=device, dtype=dtype)
1079+ grad = grad.contiguous(memory_format=channels_last)
1065 if not contig:1080 if not contig:
1066- input1 = input1[:, ::2, :, :]1081+ input1 = input1[:, ::2]
1067- grad = grad[:, ::2, :, :]1082+ grad = grad[:, ::2]
1068 input1.requires_grad_(True)1083 input1.requires_grad_(True)
1069- pool = torch.nn.AdaptiveMaxPool2d((output_height, output_width), return_indices=True).to(device)1084+ pool = mod(output_plane_size, return_indices=True).to(device)
1070 1085 
1071 ref_input = input1.detach().clone().contiguous().requires_grad_(True)1086 ref_input = input1.detach().clone().contiguous().requires_grad_(True)
1072 ref_grad = grad.detach().clone().contiguous()1087 ref_grad = grad.detach().clone().contiguous()
1073- ref_pool = torch.nn.AdaptiveMaxPool2d((output_height, output_width), return_indices=True).to(device)1088+ ref_pool = mod(output_plane_size, return_indices=True).to(device)
1074 1089 
1075 out, ind = pool(input1)1090 out, ind = pool(input1)
1076 out.backward(grad)1091 out.backward(grad)
1077 ref_out, ref_ind = ref_pool(ref_input)1092 ref_out, ref_ind = ref_pool(ref_input)
1078 ref_out.backward(ref_grad)1093 ref_out.backward(ref_grad)
1079 1094 
1080- self.assertTrue(out.is_contiguous(memory_format=torch.channels_last))1095+ # channels_last_3d case does not return channels_last_3d outputs
1096+ if n_plane_dims == 2:
1097+ self.assertTrue(out.is_contiguous(memory_format=channels_last))
1098+ self.assertTrue(ind.is_contiguous(memory_format=channels_last))
1081 self.assertTrue(ref_out.is_contiguous())1099 self.assertTrue(ref_out.is_contiguous())
1082- self.assertTrue(ind.is_contiguous(memory_format=torch.channels_last))
1083 self.assertTrue(ref_ind.is_contiguous())1100 self.assertTrue(ref_ind.is_contiguous())
1084 self.assertEqual(out, ref_out)1101 self.assertEqual(out, ref_out)
1085 self.assertEqual(ind, ref_ind)1102 self.assertEqual(ind, ref_ind)
1086 self.assertEqual(input1.grad, ref_input.grad)1103 self.assertEqual(input1.grad, ref_input.grad)
1087 1104 
1088 for contig in [True, False]:1105 for contig in [True, False]:
1089- helper(4, 8, 10, 10, 7, 7, contig)1106+ helper((4, 8, 10, 10), (7, 7), contig)
1090- helper(4, 8, 9, 14, 5, 8, contig)1107+ helper((4, 8, 9, 14), (5, 8), contig)
1091- helper(4, 8, 11, 11, 1, 1, contig)1108+ helper((4, 8, 11, 11), (1, 1), contig)
1109+ helper((2, 1, 3, 3), (1, 1), contig)
1110+ helper((4, 8, 10, 10, 10), (7, 7, 7), contig)
1111+ helper((4, 8, 11, 11, 11), (1, 1, 1), contig)
1112+ helper((2, 1, 3, 3, 3), (1, 1, 1), contig)
1092 1113 
1093 @dtypes(torch.float, torch.double)1114 @dtypes(torch.float, torch.double)
1094 def test_pooling_max_nhwc(self, device, dtype):1115 def test_pooling_max_nhwc(self, device, dtype):
@@ -336,13 +336,13 @@ class TestPruningNN(NNTestCase):
336 """336 """
337 m = nn.Linear(5, 2, bias=False)337 m = nn.Linear(5, 2, bias=False)
338 338 
339- tensor_id = id(list(m.parameters())[0])339+ tensor_id = id(next(iter(m.parameters())))
340 340 
341 prune.random_unstructured(m, name="weight", amount=0.9)341 prune.random_unstructured(m, name="weight", amount=0.9)
342- self.assertEqual(tensor_id, id(list(m.parameters())[0]))342+ self.assertEqual(tensor_id, id(next(iter(m.parameters()))))
343 343 
344 prune.remove(m, "weight")344 prune.remove(m, "weight")
345- self.assertEqual(tensor_id, id(list(m.parameters())[0]))345+ self.assertEqual(tensor_id, id(next(iter(m.parameters()))))
346 346 
347 def test_random_pruning_pickle(self):347 def test_random_pruning_pickle(self):
348 modules = [nn.Linear(5, 7), nn.Conv3d(2, 2, 2)]348 modules = [nn.Linear(5, 7), nn.Conv3d(2, 2, 2)]
@@ -35,11 +35,11 @@ import torch_npu.testing
35from torch.testing._internal.common_dtype import integral_types, get_all_math_dtypes, floating_types35from torch.testing._internal.common_dtype import integral_types, get_all_math_dtypes, floating_types
36from torch.testing._internal.common_utils import freeze_rng_state, run_tests, TestCase, skipIfNoLapack, skipIfRocm, \36from torch.testing._internal.common_utils import freeze_rng_state, run_tests, TestCase, skipIfNoLapack, skipIfRocm, \
37 TEST_NUMPY, TEST_SCIPY, TEST_WITH_CROSSREF, TEST_WITH_ROCM, \37 TEST_NUMPY, TEST_SCIPY, TEST_WITH_CROSSREF, TEST_WITH_ROCM, \
38- download_file, get_function_arglist, load_tests, skipIfMps,\38+ download_file, get_function_arglist, load_tests, skipIfMps, \
39 IS_PPC, TEST_PRIVATEUSE1, custom_device_mod, \39 IS_PPC, TEST_PRIVATEUSE1, custom_device_mod, \
40 parametrize as parametrize_test, subtest, instantiate_parametrized_tests, \40 parametrize as parametrize_test, subtest, instantiate_parametrized_tests, \
41 skipIfTorchDynamo, IS_WINDOWS, gcIfJetson, set_default_dtype41 skipIfTorchDynamo, IS_WINDOWS, gcIfJetson, set_default_dtype
42-from torch.testing._internal.common_cuda import TEST_CUDNN, TEST_CUDNN_VERSION42+from torch.testing._internal.common_cuda import TEST_CUDNN, TEST_CUDNN_VERSION, PLATFORM_SUPPORTS_FLASH_ATTENTION
43from torch.testing._internal.common_nn import NNTestCase, NewModuleTest, CriterionTest, \43from torch.testing._internal.common_nn import NNTestCase, NewModuleTest, CriterionTest, \
44 module_tests, criterion_tests, loss_reference_fns, _create_basic_net, \44 module_tests, criterion_tests, loss_reference_fns, _create_basic_net, \
45 ctcloss_reference, new_module_tests, single_batch_reference_fn, _test_bfloat16_ops, _test_module_empty_input45 ctcloss_reference, new_module_tests, single_batch_reference_fn, _test_bfloat16_ops, _test_module_empty_input
@@ -5321,6 +5321,14 @@ tensor(..., device='meta', size=(1,), requires_grad=True)""")
5321 bn.load_state_dict(empty_dict, strict=False)5321 bn.load_state_dict(empty_dict, strict=False)
5322 self.assertEqual(bn.state_dict()["num_batches_tracked"], torch.tensor(10))5322 self.assertEqual(bn.state_dict()["num_batches_tracked"], torch.tensor(10))
5323 5323 
5324+ # test that when `num_batches_tracked` is not in loaded state_dict,
5325+ # meta num_batches_tracked is still replaced with singleton 0 tensor
5326+ with torch.device('meta'):
5327+ meta_bn = torch.nn.BatchNorm2d(3)
5328+ self.assertTrue(meta_bn.num_batches_tracked.device == torch.device('meta'))
5329+ meta_bn.load_state_dict(empty_dict, assign=True, strict=False)
5330+ self.assertEqual(meta_bn.state_dict()["num_batches_tracked"], torch.tensor(0))
5331+ 
5324 def test_pairwise_distance(self):5332 def test_pairwise_distance(self):
5325 input1 = torch.randn(4, 4, requires_grad=True, dtype=torch.double)5333 input1 = torch.randn(4, 4, requires_grad=True, dtype=torch.double)
5326 input2 = torch.randn(4, 4, requires_grad=True, dtype=torch.double)5334 input2 = torch.randn(4, 4, requires_grad=True, dtype=torch.double)
@@ -5395,6 +5403,22 @@ tensor(..., device='meta', size=(1,), requires_grad=True)""")
5395 result = torch.nn.functional.cosine_embedding_loss(input1, input2, target)5403 result = torch.nn.functional.cosine_embedding_loss(input1, input2, target)
5396 self.assertEqual(result.item(), expected.item(), atol=0.001, rtol=0)5404 self.assertEqual(result.item(), expected.item(), atol=0.001, rtol=0)
5397 5405 
5406+ def test_cosine_embedding_loss_error_on_diff_shapes(self):
5407+ for device in device_():
5408+ input1 = torch.empty((0, 0), dtype=torch.double, device=device)
5409+ input2 = torch.empty((0,), dtype=torch.double, device=device)
5410+ target = torch.empty((0,), dtype=torch.int, device=device)
5411+ with self.assertRaisesRegex(RuntimeError, ".*expects 2D.*"):
5412+ torch.nn.functional.cosine_embedding_loss(input1, input2, target)
5413+ 
5414+ def test_cosine_embedding_loss_error_on_nonexpandable_shapes(self):
5415+ for device in device_():
5416+ input1 = torch.empty((1, 5), dtype=torch.double, device=device)
5417+ input2 = torch.empty((1, 6), dtype=torch.double, device=device)
5418+ target = torch.ones((1,), dtype=torch.int, device=device)
5419+ with self.assertRaisesRegex(RuntimeError, ".*must match the size.*"):
5420+ torch.nn.functional.cosine_embedding_loss(input1, input2, target)
5421+ 
5398 def test_kl_div_with_diff_type(self):5422 def test_kl_div_with_diff_type(self):
5399 for device in device_():5423 for device in device_():
5400 input1 = torch.tensor([[2, 3, 5], [3, 2, 1]], dtype=torch.double, device=device)5424 input1 = torch.tensor([[2, 3, 5], [3, 2, 1]], dtype=torch.double, device=device)
@@ -6832,7 +6856,7 @@ tensor(..., device='meta', size=(1,), requires_grad=True)""")
6832 elif weight_layout == torch.sparse_coo:6856 elif weight_layout == torch.sparse_coo:
6833 module.weight = nn.Parameter(module.weight.to_sparse_coo())6857 module.weight = nn.Parameter(module.weight.to_sparse_coo())
6834 else:6858 else:
6835- assert(0)6859+ raise AssertionError()
6836 6860 
6837 inp = torch.randn(4, requires_grad=True, device=device)6861 inp = torch.randn(4, requires_grad=True, device=device)
6838 res = module(inp)6862 res = module(inp)
@@ -7919,7 +7943,7 @@ class TestNNDeviceType(NNTestCase):
7919 mean = out_reshaped.mean(-1)7943 mean = out_reshaped.mean(-1)
7920 var = out_reshaped.var(-1, unbiased=False)7944 var = out_reshaped.var(-1, unbiased=False)
7921 7945 
7922- delta = 1e-1 if dtype == torch.bfloat16 else 1e-57946+ delta = 1e-1 if (dtype == torch.bfloat16 or dtype == torch.half) else 1e-5
7923 self.assertEqual(torch.abs(mean.data).mean(), 0, atol=delta, rtol=0)7947 self.assertEqual(torch.abs(mean.data).mean(), 0, atol=delta, rtol=0)
7924 self.assertEqual(torch.abs(var.data).mean(), 1, atol=delta, rtol=0)7948 self.assertEqual(torch.abs(var.data).mean(), 1, atol=delta, rtol=0)
7925 7949 
@@ -7953,12 +7977,12 @@ class TestNNDeviceType(NNTestCase):
7953 output.sum().backward()7977 output.sum().backward()
7954 self.assertEqualTypeString(output, input1)7978 self.assertEqualTypeString(output, input1)
7955 7979 
7956- def _test_LayerNorm_cpu_mixed_dtype(self, device):7980+ def _test_LayerNorm_cpu_mixed_dtype(self, device, dtype):
7957 for elementwise_affine in [True, False]:7981 for elementwise_affine in [True, False]:
7958- # layer norm input1 shape is normalized to m x n, cpu vectorized on n,7982+ # layer norm input shape is normalized to m x n, cpu vectorized on n,
7959 # so make sure n exceeds vector length7983 # so make sure n exceeds vector length
7960- input1 = torch.empty(2, 3, 11, 3, device=device, dtype=torch.bfloat16).random_(1, 10)7984+ input1 = torch.empty(2, 3, 11, 3, device=device, dtype=dtype).random_(1, 10)
7961- m = nn.LayerNorm([11, 3], elementwise_affine=elementwise_affine).to(device, torch.bfloat16)7985+ m = nn.LayerNorm([11, 3], elementwise_affine=elementwise_affine).to(device, dtype)
7962 7986 
7963 # fp327987 # fp32
7964 m_fp32 = deepcopy(m).to(device, torch.float)7988 m_fp32 = deepcopy(m).to(device, torch.float)
@@ -7966,21 +7990,21 @@ class TestNNDeviceType(NNTestCase):
7966 out_fp32 = m_fp32(x_fp32)7990 out_fp32 = m_fp32(x_fp32)
7967 out_fp32.sum().backward()7991 out_fp32.sum().backward()
7968 7992 
7969- # bf167993+ # bf16/half
7970 m_bf16 = deepcopy(m)7994 m_bf16 = deepcopy(m)
7971 x_bf16 = input1.clone().detach().requires_grad_()7995 x_bf16 = input1.clone().detach().requires_grad_()
7972 out_bf16 = m_bf16(x_bf16)7996 out_bf16 = m_bf16(x_bf16)
7973 out_bf16.sum().backward()7997 out_bf16.sum().backward()
7974 7998 
7975- # bf16 mixed type7999+ # bf16/half mixed type
7976 m_mix = deepcopy(m).to(device, torch.float)8000 m_mix = deepcopy(m).to(device, torch.float)
7977 x_mix = input1.clone().detach().requires_grad_()8001 x_mix = input1.clone().detach().requires_grad_()
7978 out_mix = m_mix(x_mix)8002 out_mix = m_mix(x_mix)
7979 out_mix.sum().backward()8003 out_mix.sum().backward()
7980- self.assertEqual(out_fp32.bfloat16(), out_bf16)8004+ self.assertEqual(out_fp32.to(dtype=dtype), out_bf16)
7981- self.assertEqual(out_fp32.bfloat16(), out_mix)8005+ self.assertEqual(out_fp32.to(dtype=dtype), out_mix)
7982- self.assertEqual(x_fp32.grad.bfloat16(), x_bf16.grad, atol=1e-1, rtol=1e-1)8006+ self.assertEqual(x_fp32.grad.to(dtype=dtype), x_bf16.grad, atol=1e-1, rtol=1e-1)
7983- self.assertEqual(x_fp32.grad.bfloat16(), x_mix.grad, atol=1e-1, rtol=1e-1)8007+ self.assertEqual(x_fp32.grad.to(dtype=dtype), x_mix.grad, atol=1e-1, rtol=1e-1)
7984 8008 
7985 def _test_GroupNorm_general(self, device, dtype=torch.float):8009 def _test_GroupNorm_general(self, device, dtype=torch.float):
7986 good_shape_g = {8010 good_shape_g = {
@@ -8245,6 +8269,21 @@ class TestNNDeviceType(NNTestCase):
8245 8269 
8246 self.assertEqual(scipy_ary, gridsample_ary.reshape_as(scipy_ary))8270 self.assertEqual(scipy_ary, gridsample_ary.reshape_as(scipy_ary))
8247 8271 
8272+ @onlyPRIVATEUSE1
8273+ @largeTensorTest("60GB", "cpu")
8274+ @largeTensorTest("16GB", "npu")
8275+ def test_avg_pool_large_tensor(self, device):
8276+ # test for pytorch issues 113833
8277+ a = torch.randn(128, 256, 256, 256, dtype=torch.half, device=device, requires_grad=True)
8278+ a_cpu = a.detach().cpu().float()
8279+ m = torch.nn.AvgPool2d(2)
8280+ out = m(a)
8281+ a_cpu.requires_grad = True
8282+ out.sum().backward()
8283+ o_cpu = m(a_cpu)
8284+ o_cpu.sum().backward()
8285+ self.assertTrue(torch.allclose(a.grad.cpu(), a_cpu.grad.half()))
8286+ 
8248 @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'),8287 @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'),
8249 "Scipy v1.0 and/or numpy not found")8288 "Scipy v1.0 and/or numpy not found")
8250 @tf32_on_and_off(0.005)8289 @tf32_on_and_off(0.005)
@@ -8471,14 +8510,16 @@ class TestNNDeviceType(NNTestCase):
8471 def test_LayerNorm_general(self, device):8510 def test_LayerNorm_general(self, device):
8472 self._test_LayerNorm_general(device)8511 self._test_LayerNorm_general(device)
8473 8512 
8474- if self.device_type == 'cpu' or self.device_type == torch._C._get_privateuse1_backend_name():8513+ if self.device_type == 'npu' or self.device_type == 'cpu':
8475- self._test_LayerNorm_general(device, dtype=torch.bfloat16)8514+ for dtype in [torch.half, torch.bfloat16]:
8515+ self._test_LayerNorm_general(device, dtype=dtype)
8476 8516 
8477- if self.device_type == torch._C._get_privateuse1_backend_name():8517+ if self.device_type == 'npu':
8478 self._test_LayerNorm_cuda_half(device)8518 self._test_LayerNorm_cuda_half(device)
8479 8519 
8480 if self.device_type == 'cpu':8520 if self.device_type == 'cpu':
8481- self._test_LayerNorm_cpu_mixed_dtype(device)8521+ for dtype in [torch.half, torch.bfloat16]:
8522+ self._test_LayerNorm_cpu_mixed_dtype(device, dtype=dtype)
8482 8523 
8483 @onlyNativeDeviceTypes8524 @onlyNativeDeviceTypes
8484 def test_LayerNorm_numeric(self, device):8525 def test_LayerNorm_numeric(self, device):
@@ -11523,6 +11564,19 @@ class TestNNDeviceType(NNTestCase):
11523 with torch.no_grad():11564 with torch.no_grad():
11524 self.assertTrue(torch.allclose(input1.grad.cpu(), input_cpu.grad, rtol=rtol, atol=atol))11565 self.assertTrue(torch.allclose(input1.grad.cpu(), input_cpu.grad, rtol=rtol, atol=atol))
11525 11566 
11567+ # Ref: pytorch issue 108345
11568+ @onlyPRIVATEUSE1
11569+ @largeTensorTest("20GB", "cpu")
11570+ @largeTensorTest("20GB", "npu")
11571+ @parametrize_test("reduction", ("none", "mean", "sum"))
11572+ def test_cross_entropy_64bit(self, device, reduction):
11573+ labels = torch.zeros(190, 50, dtype=torch.long, device=device)
11574+ logits = torch.ones(190, 229000, 50, dtype=torch.float, device=device)
11575+ loss = torch.nn.functional.cross_entropy(logits, labels)
11576+ loss_cpu = torch.nn.functional.cross_entropy(logits.cpu(), labels.cpu())
11577+ print(logits.numel(), labels.numel(), loss.numel())
11578+ self.assertTrue(torch.allclose(loss_cpu, loss.cpu(), rtol=1e-4, atol=1e-4))
11579+ 
11526 def _nll_loss_helper(self, input_size, reduction, expected, device):11580 def _nll_loss_helper(self, input_size, reduction, expected, device):
11527 input1 = torch.rand(input_size, requires_grad=True, device=device)11581 input1 = torch.rand(input_size, requires_grad=True, device=device)
11528 num_channels = input_size[1]11582 num_channels = input_size[1]
@@ -12321,6 +12375,8 @@ class TestNNDeviceType(NNTestCase):
12321 @dtypes(torch.float)12375 @dtypes(torch.float)
12322 @dtypesIfPRIVATEUSE1(torch.double, torch.float, torch.half)12376 @dtypesIfPRIVATEUSE1(torch.double, torch.float, torch.half)
12323 def test_transformerencoderlayer(self, device, dtype):12377 def test_transformerencoderlayer(self, device, dtype):
12378+ if TEST_WITH_ROCM and PLATFORM_SUPPORTS_FLASH_ATTENTION and dtype == torch.half:
12379+ self.skipTest("Skip on ROCM due to Flash Attention tolerances")
12324 # this is a deterministic test for TransformerEncoderLayer12380 # this is a deterministic test for TransformerEncoderLayer
12325 d_model = 412381 d_model = 4
12326 nhead = 212382 nhead = 2
@@ -12531,6 +12587,8 @@ class TestNNDeviceType(NNTestCase):
12531 @dtypes(torch.float)12587 @dtypes(torch.float)
12532 @dtypesIfPRIVATEUSE1(torch.half, torch.float)12588 @dtypesIfPRIVATEUSE1(torch.half, torch.float)
12533 def test_transformerencoderlayer_gelu(self, device, dtype):12589 def test_transformerencoderlayer_gelu(self, device, dtype):
12590+ if TEST_WITH_ROCM and PLATFORM_SUPPORTS_FLASH_ATTENTION and dtype == torch.half:
12591+ self.skipTest("Skip on ROCM due to Flash Attention tolerances")
12534 # this is a deterministic test for TransformerEncoderLayer with gelu activation12592 # this is a deterministic test for TransformerEncoderLayer with gelu activation
12535 d_model = 412593 d_model = 4
12536 nhead = 212594 nhead = 2
@@ -12696,6 +12754,21 @@ class TestNNDeviceType(NNTestCase):
12696 clip_grad_norm_([p2], max_norm, norm_type=norm_type, foreach=foreach)12754 clip_grad_norm_([p2], max_norm, norm_type=norm_type, foreach=foreach)
12697 self.assertEqual(p1.grad, p2.grad)12755 self.assertEqual(p1.grad, p2.grad)
12698 12756 
12757+ # reference issue: pytorch issues 111484
12758+ @onlyPRIVATEUSE1
12759+ @largeTensorTest("41GB" if TEST_WITH_ROCM else "30GB", "cuda")
12760+ def test_softmax_forward_64bit_indexing(self, device):
12761+ batch_size = 70
12762+ seq_len = 2048
12763+ vocab_size = 50000
12764+ 
12765+ shift_labels = torch.zeros(batch_size, seq_len - 1, dtype=torch.long, device=device)
12766+ logits = torch.ones(batch_size, seq_len - 1, vocab_size, dtype=torch.float16, device=device)
12767+ loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
12768+ nll = loss_fct(logits.permute(0, 2, 1), shift_labels).float()
12769+ rtol, atol = torch.testing._comparison.get_tolerances(torch.float16, rtol=None, atol=None)
12770+ self.assertEqual(nll, torch.ones_like(nll) * torch.log(torch.tensor(vocab_size)), rtol=rtol, atol=atol)
12771+ 
12699 @onlyPRIVATEUSE112772 @onlyPRIVATEUSE1
12700 @largeTensorTest("20GB", "npu")12773 @largeTensorTest("20GB", "npu")
12701 def test_softmax_backward_64bit_indexing(self, device):12774 def test_softmax_backward_64bit_indexing(self, device):
@@ -12705,6 +12778,21 @@ class TestNNDeviceType(NNTestCase):
12705 out = torch._softmax_backward_data(x, x, 2, x.dtype)12778 out = torch._softmax_backward_data(x, x, 2, x.dtype)
12706 self.assertEqual(out[0, 0, 0], 1 / numel)12779 self.assertEqual(out[0, 0, 0], 1 / numel)
12707 12780 
12781+ # reference issue: pytorch issues 68248
12782+ @onlyPRIVATEUSE1
12783+ def test_adaptiveavg_pool1d_shmem(self, device):
12784+ x = torch.randn(1, 256, 1, 5000, device=device).to(memory_format=torch.channels_last)
12785+ x_cpu = x.cpu()
12786+ x_cpu.requires_grad_()
12787+ x.requires_grad_()
12788+ y = torch.nn.functional.adaptive_avg_pool2d(x, (1, 256))
12789+ y_cpu = torch.nn.functional.adaptive_avg_pool2d(x_cpu, (1, 256))
12790+ grad = torch.randn_like(y)
12791+ grad_cpu = grad.cpu()
12792+ y.backward(grad)
12793+ y_cpu.backward(grad_cpu)
12794+ self.assertEqual(x.grad, x_cpu.grad)
12795+ 
12708 @skipMeta12796 @skipMeta
12709 def test_channel_shuffle(self, device):12797 def test_channel_shuffle(self, device):
12710 # 3D tensor12798 # 3D tensor
@@ -67,6 +67,7 @@ if TEST_PRIVATEUSE1:
67_cycles_per_ms = None67_cycles_per_ms = None
68 68 
69 69 
70+@torch.testing._internal.common_utils.markDynamoStrictTest
70class TestNpu(TestCase):71class TestNpu(TestCase):
71 _do_cuda_memory_leak_check = True72 _do_cuda_memory_leak_check = True
72 _do_cuda_non_default_stream = True73 _do_cuda_non_default_stream = True
@@ -80,6 +81,28 @@ class TestNpu(TestCase):
80 del self.autocast_lists81 del self.autocast_lists
81 super().tearDown()82 super().tearDown()
82 83 
84+ def test_pinned_memory_with_cudaregister(self):
85+ torch_npu.npu.memory._set_allocator_settings("pinned_use_npu_host_register:True,pinned_num_register_threads:8")
86+ t = torch.ones(20)
87+ self.assertFalse(t.is_pinned())
88+ try:
89+ pinned_t = torch.ones(1 << 21).pin_memory()
90+ self.assertTrue(pinned_t.is_pinned())
91+ pinned_t = torch.ones(1 << 24).pin_memory()
92+ self.assertTrue(pinned_t.is_pinned())
93+ except RuntimeError as e:
94+ # Some NPUs don't support same address space on host and device side
95+ pass
96+ 
97+ def test_pinned_memory_with_cudaregister_multithread(self):
98+ num_threads = 4
99+ threads = [threading.Thread(target=self.test_pinned_memory_with_cudaregister)
100+ for t in range(num_threads)]
101+ for thread in threads:
102+ thread.start()
103+ for thread in threads:
104+ thread.join()
105+ 
83 def test_cudart_register(self):106 def test_cudart_register(self):
84 t = torch.ones(20)107 t = torch.ones(20)
85 self.assertFalse(t.is_pinned())108 self.assertFalse(t.is_pinned())
@@ -854,7 +877,7 @@ except RuntimeError as e:
854 @staticmethod877 @staticmethod
855 def backward(ctx, grad):878 def backward(ctx, grad):
856 self.assertEqual(torch_npu.npu.current_stream(), ctx.stream)879 self.assertEqual(torch_npu.npu.current_stream(), ctx.stream)
857- # delays the operation in the the background stream880+ # delays the operation in the background stream
858 torch_npu.npu._sleep(1000 * 5000)881 torch_npu.npu._sleep(1000 * 5000)
859 return grad * ctx.val, None882 return grad * ctx.val, None
860 883 
@@ -1361,7 +1384,7 @@ torch_npu.npu.synchronize()
1361 scaler.scale(z).backward()1384 scaler.scale(z).backward()
1362 scaler.step(optimizer)1385 scaler.step(optimizer)
1363 scaler.update()1386 scaler.update()
1364- assert(scaler._scale != float('inf') and scaler._scale != float('nan'))1387+ assert scaler._scale != float('inf') and scaler._scale != float('nan')
1365 1388 
1366 def test_grad_scaling_clipping(self):1389 def test_grad_scaling_clipping(self):
1367 def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):1390 def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
@@ -1811,8 +1834,7 @@ torch_npu.npu.synchronize()
1811 ('TORCH_CUDNN_V8_API_DISABLED' in os.environ and1834 ('TORCH_CUDNN_V8_API_DISABLED' in os.environ and
1812 int(os.environ['TORCH_CUDNN_V8_API_DISABLED']) or1835 int(os.environ['TORCH_CUDNN_V8_API_DISABLED']) or
1813 torch_npu.npu.get_device_capability() < (8, 0))1836 torch_npu.npu.get_device_capability() < (8, 0))
1814- should_error_from_not_implemented = should_error_from_cudnn or 'thnn' in op \1837+ should_error_from_not_implemented = should_error_from_cudnn
1815- or 'fused' in op or 'gru' in op or op == '_thnn_fused_lstm_cell' or op == 'lstm_cell'
1816 if not skip_test:1838 if not skip_test:
1817 if should_error_from_not_implemented:1839 if should_error_from_not_implemented:
1818 with self.assertRaises(RuntimeError, msg=str(op) + ' should not be supported for bfloat16!'):1840 with self.assertRaises(RuntimeError, msg=str(op) + ' should not be supported for bfloat16!'):
@@ -2137,7 +2159,7 @@ torch_npu.npu.synchronize()
2137 def test_graph_is_current_stream_capturing(self):2159 def test_graph_is_current_stream_capturing(self):
2138 self.assertFalse(torch_npu.npu.is_current_stream_capturing())2160 self.assertFalse(torch_npu.npu.is_current_stream_capturing())
2139 2161 
2140- if (TEST_PRIVATEUSE1 and (not TEST_WITH_ROCM) and int(torch.version.npu.split(".")[0]) >= 11):2162+ if (TEST_PRIVATEUSE1 and (not TEST_WITH_ROCM)):
2141 s = torch_npu.npu.Stream()2163 s = torch_npu.npu.Stream()
2142 with torch_npu.npu.stream(s):2164 with torch_npu.npu.stream(s):
2143 g = torch_npu.npu.NPUGraph()2165 g = torch_npu.npu.NPUGraph()
@@ -2165,6 +2187,41 @@ torch_npu.npu.synchronize()
2165 2187 
2166 self.assertTrue(b.sum().item() == 11000.)2188 self.assertTrue(b.sum().item() == 11000.)
2167 2189 
2190+ @unittest.skipIf(not TEST_PRIVATEUSE1, "NPU not available")
2191+ def test_graph_capture_reset_recapture(self):
2192+ s = torch_npu.npu.Stream()
2193+ 
2194+ with torch_npu.npu.stream(s):
2195+ a = torch.full((1000,), 1, device="npu")
2196+ g = torch_npu.npu.NPUGraph()
2197+ torch_npu.npu.empty_cache()
2198+ g.capture_begin()
2199+ b = a
2200+ for _ in range(10):
2201+ b = b + 1
2202+ g.capture_end()
2203+ torch_npu.npu.current_stream().wait_stream(s)
2204+ 
2205+ g.replay()
2206+ 
2207+ self.assertTrue(b.sum().item() == 11000.)
2208+ 
2209+ g.reset()
2210+ 
2211+ with torch_npu.npu.stream(s):
2212+ g.capture_begin()
2213+ b.fill_(2.0)
2214+ for _ in range(10):
2215+ b = b + 2
2216+ g.capture_end()
2217+ torch_npu.npu.current_stream().wait_stream(s)
2218+ 
2219+ g.replay()
2220+ self.assertTrue(b.sum().item() == 22000.)
2221+ 
2222+ g.reset()
2223+ del g
2224+ 
2168 @unittest.skipIf(not TEST_PRIVATEUSE1, "NPU not available")2225 @unittest.skipIf(not TEST_PRIVATEUSE1, "NPU not available")
2169 def test_graph_error(self):2226 def test_graph_error(self):
2170 # We need to run this test in a separate thread as the error we trigger2227 # We need to run this test in a separate thread as the error we trigger
@@ -3041,6 +3098,9 @@ exit(2)
3041 ] + [3098 ] + [
3042 (optimizer_ctor, {"lr": 0.1, "betas": (0.8, 0.7), "fused": True, "amsgrad": amsgrad})3099 (optimizer_ctor, {"lr": 0.1, "betas": (0.8, 0.7), "fused": True, "amsgrad": amsgrad})
3043 for optimizer_ctor, amsgrad in product((torch.optim.Adam, torch.optim.AdamW), (False, True))3100 for optimizer_ctor, amsgrad in product((torch.optim.Adam, torch.optim.AdamW), (False, True))
3101+ ] + [
3102+ (torch.optim.ASGD, {"lr": 0.1, "foreach": True, "maximize": maximize, "weight_decay": weight_decay})
3103+ for maximize, weight_decay in product((False, True), (0.0, 0.1))
3044 ]3104 ]
3045 3105 
3046 for optimizer_ctor, kwargs in cases:3106 for optimizer_ctor, kwargs in cases:
@@ -3222,6 +3282,30 @@ exit(2)
3222 # Exception would Corrupt Process and make other tests fail3282 # Exception would Corrupt Process and make other tests fail
3223 # self.assertTrue(throws_on_cuda_event("global"))3283 # self.assertTrue(throws_on_cuda_event("global"))
3224 3284 
3285+ @unittest.skipIf(not TEST_PRIVATEUSE1, "NPU not available")
3286+ def test_cuda_graph_allocator_propagates_stream(self):
3287+ segments = torch_npu.npu.memory_snapshot()
3288+ existing_pools = {s["segment_pool_id"] for s in segments}
3289+ x = torch.randn(10240000, device="npu")
3290+ y = torch.rand_like(x)
3291+ g = torch_npu.npu.NPUGraph()
3292+ s0 = torch_npu.npu.Stream()
3293+ s1 = torch_npu.npu.Stream()
3294+ s0.wait_stream(torch_npu.npu.current_stream())
3295+ with torch_npu.npu.stream(s0):
3296+ g.capture_begin()
3297+ z = x + y
3298+ with torch_npu.npu.stream(s1):
3299+ s1.wait_stream(s0)
3300+ w = z + y
3301+ s0.wait_stream(s1)
3302+ with torch_npu.npu.stream(s0):
3303+ g.capture_end()
3304+ segments = torch_npu.npu.memory_snapshot()
3305+ x = [s["segment_pool_id"] for s in segments if s["segment_pool_id"] not in existing_pools]
3306+ self.assertEqual(len(x), 2)
3307+ self.assertEqual(x[0], x[1])
3308+ 
3225 def test_batch_norm_gather_stats(self):3309 def test_batch_norm_gather_stats(self):
3226 input1 = torch.randn(1, 3, 3, 3, device='npu')3310 input1 = torch.randn(1, 3, 3, 3, device='npu')
3227 mean, invstd = torch.batch_norm_gather_stats(3311 mean, invstd = torch.batch_norm_gather_stats(
@@ -3309,14 +3393,24 @@ exit(2)
3309 with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"):3393 with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"):
3310 torch.addmm(s, m1, m2)3394 torch.addmm(s, m1, m2)
3311 3395 
3312- @unittest.skipIf(not TEST_MULTINPU, "Testing on one NPU is sufficient")3396+ @unittest.skipIf(TEST_MULTINPU, "Testing on one NPU is sufficient")
3313 def test_lazy_init(self):3397 def test_lazy_init(self):
3314 """ Validate that no NPU calls are made during `import torch` call"""3398 """ Validate that no NPU calls are made during `import torch` call"""
3315- from subprocess import check_output3399+ def check_output(script: str) -> str:
3400+ return subprocess.check_output([sys.executable, "-c", script]).decode("ascii").strip()
3401+ 
3316 VISIBLE_DEVICES = "HIP_VISIBLE_DEVICES" if TEST_WITH_ROCM else "CUDA_VISIBLE_DEVICES"3402 VISIBLE_DEVICES = "HIP_VISIBLE_DEVICES" if TEST_WITH_ROCM else "CUDA_VISIBLE_DEVICES"
3317- test_script = f"import os; import torch; import torch_npu; os.environ['{VISIBLE_DEVICES}']='32';print(torch_npu.npu.device_count())"3403+ test_script = f"import os; import torch;os.environ['{VISIBLE_DEVICES}']='32';print(torch.cuda.device_count())"
3318- rc = check_output([sys.executable, '-c', test_script]).decode("ascii").strip()3404+ rc = check_output(test_script)
3319 self.assertEqual(rc, "0")3405 self.assertEqual(rc, "0")
3406+ if not TEST_WITH_ROCM:
3407+ # Check that `cuInit` was not called during the import
3408+ # By using ctypes and calling cuDeviceCountGet() and expect CUDA_ERROR_NOT_INITIALIZED == 3
3409+ # See pytorch issues 116276 for more details
3410+ libcuda_name = "libcuda.so.1" if not IS_WINDOWS else "nvcuda.dll"
3411+ cuda_driver_api_call = f"ctypes.CDLL('{libcuda_name}').cuDeviceGetCount(ctypes.byref(x))"
3412+ rc = check_output(f"import torch; import ctypes;x=ctypes.c_int(-1);print({cuda_driver_api_call})")
3413+ self.assertEqual(rc, "3")
3320 3414 
3321 3415 
3322class TestNpuMallocAsync(TestCase):3416class TestNpuMallocAsync(TestCase):
@@ -3595,9 +3689,16 @@ class TestNpuMallocAsync(TestCase):
3595 3689 
3596 pow2_div2_mem = torch_npu.npu.memory_stats()[key_allocated]3690 pow2_div2_mem = torch_npu.npu.memory_stats()[key_allocated]
3597 if not TEST_NPUMALLOCASYNC:3691 if not TEST_NPUMALLOCASYNC:
3598- # not supported with the cudaMallocAsync backend3692+ # not supported with the npuMallocAsync backend
3599 self.assertTrue(pow2_div2_mem - start_mem == power2_div(nbytes_big, 2))3693 self.assertTrue(pow2_div2_mem - start_mem == power2_div(nbytes_big, 2))
3600 3694 
3695+ torch_npu.npu.memory.empty_cache()
3696+ torch_npu.npu.memory._set_allocator_settings("release_lock_on_npumalloc:True")
3697+ start_mem = torch_npu.npu.memory_stats()[key_allocated]
3698+ w = torch.rand(nelems, device='npu')
3699+ reg_mem = torch_npu.npu.memory_stats()[key_allocated]
3700+ self.assertTrue(reg_mem - start_mem == nbytes)
3701+ 
3601 with self.assertRaises(RuntimeError):3702 with self.assertRaises(RuntimeError):
3602 torch_npu.npu.memory._set_allocator_settings("foo:1,bar:2")3703 torch_npu.npu.memory._set_allocator_settings("foo:1,bar:2")
3603 3704 
@@ -3607,6 +3708,18 @@ class TestNpuMallocAsync(TestCase):
3607 with self.assertRaises(RuntimeError):3708 with self.assertRaises(RuntimeError):
3608 torch_npu.npu.memory._set_allocator_settings("max_split_size_mb:2")3709 torch_npu.npu.memory._set_allocator_settings("max_split_size_mb:2")
3609 3710 
3711+ with self.assertRaises(RuntimeError):
3712+ torch_npu.npu.memory._set_allocator_settings("release_lock_on_npumalloc:none")
3713+ 
3714+ with self.assertRaises(RuntimeError):
3715+ torch_npu.npu.memory._set_allocator_settings("pinned_use_npu_host_register:none")
3716+ 
3717+ with self.assertRaises(RuntimeError):
3718+ torch_npu.npu.memory._set_allocator_settings("pinned_num_register_threads:none")
3719+ 
3720+ with self.assertRaises(RuntimeError):
3721+ torch_npu.npu.memory._set_allocator_settings("pinned_num_register_threads:1024")
3722+ 
3610 def test_raises_oom(self):3723 def test_raises_oom(self):
3611 with self.assertRaises(torch_npu.npu.OutOfMemoryError):3724 with self.assertRaises(torch_npu.npu.OutOfMemoryError):
3612 torch.empty(1024 * 1024 * 1024 * 1024, device='npu')3725 torch.empty(1024 * 1024 * 1024 * 1024, device='npu')
@@ -3798,6 +3911,7 @@ def reconstruct_from_tensor_metadata(metadata):
3798 3911 
3799 3912 
3800@unittest.skipIf(TEST_NPUMALLOCASYNC or TEST_WITH_ROCM, "NYI")3913@unittest.skipIf(TEST_NPUMALLOCASYNC or TEST_WITH_ROCM, "NYI")
3914+@torch.testing._internal.common_utils.markDynamoStrictTest
3801class TestBlockStateAbsorption(TestCase):3915class TestBlockStateAbsorption(TestCase):
3802 3916 
3803 def checkCheckpointedBlock(self, before_block, after_block):3917 def checkCheckpointedBlock(self, before_block, after_block):
@@ -133,6 +133,8 @@
133 "test_view_copy_privateuse1 (__main__.TestAutogradMultipleDispatchPRIVATEUSE1)": ["", [""]],133 "test_view_copy_privateuse1 (__main__.TestAutogradMultipleDispatchPRIVATEUSE1)": ["", [""]],
134 "test_where_functional_npu (__main__.TestAutogradDeviceTypePRIVATEUSE1)": ["", [""]],134 "test_where_functional_npu (__main__.TestAutogradDeviceTypePRIVATEUSE1)": ["", [""]],
135 "test_where_scalar_npu (__main__.TestAutogradDeviceTypePRIVATEUSE1)": ["", [""]],135 "test_where_scalar_npu (__main__.TestAutogradDeviceTypePRIVATEUSE1)": ["", [""]],
136+ "test_full_backward_pre_hooks_named_tuple_True (__main__.TestModuleHooks)": ["", [""]],
137+ "test_full_backward_pre_hooks_named_tuple_False (__main__.TestModuleHooks)": ["", [""]],
136 "test_check_inplace_nn_CELU_npu_float64 (__main__.TestModulePRIVATEUSE1)": ["", [""]],138 "test_check_inplace_nn_CELU_npu_float64 (__main__.TestModulePRIVATEUSE1)": ["", [""]],
137 "test_check_inplace_nn_ELU_npu_float64 (__main__.TestModulePRIVATEUSE1)": ["", [""]],139 "test_check_inplace_nn_ELU_npu_float64 (__main__.TestModulePRIVATEUSE1)": ["", [""]],
138 "test_check_inplace_nn_Hardswish_npu_float64 (__main__.TestModulePRIVATEUSE1)": ["", [""]],140 "test_check_inplace_nn_Hardswish_npu_float64 (__main__.TestModulePRIVATEUSE1)": ["", [""]],
@@ -2050,6 +2052,7 @@
2050 "test_CTCLoss_lengthchecks_cuda (__main__.TestNN)": ["", [""]],2052 "test_CTCLoss_lengthchecks_cuda (__main__.TestNN)": ["", [""]],
2051 "test_CTCLoss_long_targets (__main__.TestNN)": ["", [""]],2053 "test_CTCLoss_long_targets (__main__.TestNN)": ["", [""]],
2052 "test_CTCLoss_zero_infinity (__main__.TestNN)": ["", [""]],2054 "test_CTCLoss_zero_infinity (__main__.TestNN)": ["", [""]],
2055+ "test_cudnn_rnn_dropout_states_device (__main__.TestNN)": ["", [""]],
2053 "test_partial_flat_weights (__main__.TestNN)": ["", [""]],2056 "test_partial_flat_weights (__main__.TestNN)": ["", [""]],
2054 "test_batchnorm_cudnn_nhwc (__main__.TestNN)": ["", [""]],2057 "test_batchnorm_cudnn_nhwc (__main__.TestNN)": ["", [""]],
2055 "test_batchnorm_nhwc_cuda (__main__.TestNN)": ["", [""]],2058 "test_batchnorm_nhwc_cuda (__main__.TestNN)": ["", [""]],
@@ -2100,6 +2103,7 @@
2100 "test_transformerencoderlayer_gelu_npu_float16 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2103 "test_transformerencoderlayer_gelu_npu_float16 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2101 "test_transformerencoderlayer_npu_float64 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2104 "test_transformerencoderlayer_npu_float64 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2102 "test_activations_bfloat16_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2105 "test_activations_bfloat16_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2106+ "test_adaptiveavg_pool1d_shmem_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2103 "test_batchnorm_affine_npu_bfloat16 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2107 "test_batchnorm_affine_npu_bfloat16 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2104 "test_to (__main__.PackedSequenceTest)": ["", [""]],2108 "test_to (__main__.PackedSequenceTest)": ["", [""]],
2105 "test_new_spectral_norm (__main__.TestNNParametrization)": ["", [""]],2109 "test_new_spectral_norm (__main__.TestNNParametrization)": ["", [""]],
@@ -2107,6 +2111,7 @@
2107 "test_adaptive_avg_pooling_nhwc_overflow (__main__.TestPoolingNN)": ["", [""]],2111 "test_adaptive_avg_pooling_nhwc_overflow (__main__.TestPoolingNN)": ["", [""]],
2108 "test_adaptive_pooling_avg_nhwc (__main__.TestPoolingNN)": ["", [""]],2112 "test_adaptive_pooling_avg_nhwc (__main__.TestPoolingNN)": ["", [""]],
2109 "test_adaptive_pooling_avg_nhwc_non_contiguous (__main__.TestPoolingNN)": ["", [""]],2113 "test_adaptive_pooling_avg_nhwc_non_contiguous (__main__.TestPoolingNN)": ["", [""]],
2114+ "test_quantized_max_pool1d_empty_kernel (__main__.TestPoolingNN)": ["", [""]],
2110 "test_avg_pool2d_nhwc_npu_float16 (__main__.TestPoolingNNDeviceTypePRIVATEUSE1)": ["", [""]],2115 "test_avg_pool2d_nhwc_npu_float16 (__main__.TestPoolingNNDeviceTypePRIVATEUSE1)": ["", [""]],
2111 "test_adaptive_pooling_empty_output_size_npu_bfloat16 (__main__.TestPoolingNNDeviceTypePRIVATEUSE1)": ["", [""]],2116 "test_adaptive_pooling_empty_output_size_npu_bfloat16 (__main__.TestPoolingNNDeviceTypePRIVATEUSE1)": ["", [""]],
2112 "test_adaptive_pooling_empty_output_size_npu_float16 (__main__.TestPoolingNNDeviceTypePRIVATEUSE1)": ["", [""]],2117 "test_adaptive_pooling_empty_output_size_npu_float16 (__main__.TestPoolingNNDeviceTypePRIVATEUSE1)": ["", [""]],
@@ -2145,6 +2150,7 @@
2145 "test_conv_empty_input_npu_float64 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2150 "test_conv_empty_input_npu_float64 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2146 "test_conv_empty_input_npu_complex128 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2151 "test_conv_empty_input_npu_complex128 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2147 "test_LayerNorm_general_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2152 "test_LayerNorm_general_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2153+ "test_LayerNorm_general_cpu (__main__.TestNNDeviceTypeCPU)": ["", [""]],
2148 "test_lstmcell_backward_only_one_output_grad_npu_float64 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2154 "test_lstmcell_backward_only_one_output_grad_npu_float64 (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2149 "test_GroupNorm_general_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2155 "test_GroupNorm_general_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
2150 "test_GroupNorm_memory_format_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],2156 "test_GroupNorm_memory_format_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
@@ -2784,6 +2790,8 @@
2784 "test_record_stream (__main__.TestNpu)": ["", [""]],2790 "test_record_stream (__main__.TestNpu)": ["", [""]],
2785 "test_record_stream_on_shifted_view (__main__.TestNpu)": ["", [""]],2791 "test_record_stream_on_shifted_view (__main__.TestNpu)": ["", [""]],
2786 "test_caching_pinned_memory (__main__.TestNpu)": ["", [""]],2792 "test_caching_pinned_memory (__main__.TestNpu)": ["", [""]],
2793+ "test_cuda_graph_allocator_propagates_stream (__main__.TestNpu)": ["", [""]],
2794+ "test_pinned_memory_with_cudaregister (__main__.TestNpu)": ["", [""]],
2787 "test_prod_large (__main__.TestNpu)": ["", [""]],2795 "test_prod_large (__main__.TestNpu)": ["", [""]],
2788 "test_multinomial_ext (__main__.TestNpu)": ["", [""]],2796 "test_multinomial_ext (__main__.TestNpu)": ["", [""]],
2789 "test_min_max_inits (__main__.TestNpu)": ["", [""]],2797 "test_min_max_inits (__main__.TestNpu)": ["", [""]],
@@ -2826,6 +2834,7 @@
2826 "test_autocast_rnn (__main__.TestNpu)": ["", [""]],2834 "test_autocast_rnn (__main__.TestNpu)": ["", [""]],
2827 "test_graph_is_current_stream_capturing (__main__.TestNpu)": ["", [""]],2835 "test_graph_is_current_stream_capturing (__main__.TestNpu)": ["", [""]],
2828 "test_graph_capture_simple (__main__.TestNpu)": ["", [""]],2836 "test_graph_capture_simple (__main__.TestNpu)": ["", [""]],
2837+ "test_graph_capture_reset_recapture (__main__.TestNpu)": ["", [""]],
2829 "test_graph_error (__main__.TestNpu)": ["", [""]],2838 "test_graph_error (__main__.TestNpu)": ["", [""]],
2830 "test_graph_warn_if_has_zero_nodes (__main__.TestNpu)": ["", [""]],2839 "test_graph_warn_if_has_zero_nodes (__main__.TestNpu)": ["", [""]],
2831 "test_graph_capture_oom (__main__.TestNpu)": ["", [""]],2840 "test_graph_capture_oom (__main__.TestNpu)": ["", [""]],