已合并
[test][v2.7.1] add NPU validation cases for torch._functorch.config.patch and _add_batch_dim (issues #2687 #2688) #41382
ggg_0963创建于 7月12日
[test][v2.7.1] add NPU validation cases for torch._functorch.config.patch and _add_batch_dim (issues #2687 #2688) #41382
已合并
共 2 个文件变更+286-0
| @@ -0,0 +1,102 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved. | ||
| 2 | +# | ||
| 3 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 4 | +# you may not use this file except in compliance with the License. | ||
| 5 | +# You may obtain a copy of the License at | ||
| 6 | +# | ||
| 7 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 8 | +# | ||
| 9 | +# Unless required by applicable law or agreed to in writing, software | ||
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +# See the License for the specific language governing permissions and | ||
| 13 | +# limitations under the License. | ||
| 14 | + | ||
| 15 | +""" | ||
| 16 | +Add validation cases for torch._functorch.config.patch API on NPU: | ||
| 17 | +1. PyTorch community lacks sufficient and direct API validations for this API, so this file is added. | ||
| 18 | +2. This file validates torch._functorch.config.patch (extendable). | ||
| 19 | +""" | ||
| 20 | +import torch | ||
| 21 | +import torch._functorch.config as config | ||
| 22 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 23 | + | ||
| 24 | +device_type = ( | ||
| 25 | + acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 26 | +) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +class TestFunctorchConfigPatch(TestCase): | ||
| 30 | + def test_basic_patch(self): | ||
| 31 | + """Test basic config.patch context manager""" | ||
| 32 | + original_value = config.debug_assert | ||
| 33 | + with config.patch("debug_assert", not original_value): | ||
| 34 | + self.assertEqual(config.debug_assert, not original_value) | ||
| 35 | + self.assertEqual(config.debug_assert, original_value) | ||
| 36 | + | ||
| 37 | + def test_patch_dict(self): | ||
| 38 | + """Test config.patch with dict argument""" | ||
| 39 | + original_debug = config.debug_assert | ||
| 40 | + original_cse = config.cse | ||
| 41 | + patches = { | ||
| 42 | + "debug_assert": not original_debug, | ||
| 43 | + "cse": not original_cse, | ||
| 44 | + } | ||
| 45 | + with config.patch(patches): | ||
| 46 | + self.assertEqual(config.debug_assert, not original_debug) | ||
| 47 | + self.assertEqual(config.cse, not original_cse) | ||
| 48 | + self.assertEqual(config.debug_assert, original_debug) | ||
| 49 | + self.assertEqual(config.cse, original_cse) | ||
| 50 | + | ||
| 51 | + def test_patch_restore_after_exception(self): | ||
| 52 | + """Test that config is restored even after an exception""" | ||
| 53 | + original_value = config.debug_assert | ||
| 54 | + with self.assertRaises(RuntimeError): | ||
| 55 | + with config.patch("debug_assert", not original_value): | ||
| 56 | + self.assertEqual(config.debug_assert, not original_value) | ||
| 57 | + raise RuntimeError("test exception") | ||
| 58 | + self.assertEqual(config.debug_assert, original_value) | ||
| 59 | + | ||
| 60 | + def test_patch_nested(self): | ||
| 61 | + """Test nested config.patch contexts""" | ||
| 62 | + original_value = config.debug_assert | ||
| 63 | + # First level | ||
| 64 | + with config.patch("debug_assert", True): | ||
| 65 | + self.assertTrue(config.debug_assert) | ||
| 66 | + # Second level (nested) | ||
| 67 | + with config.patch("debug_assert", False): | ||
| 68 | + self.assertFalse(config.debug_assert) | ||
| 69 | + # Back to first level | ||
| 70 | + self.assertTrue(config.debug_assert) | ||
| 71 | + # Back to original | ||
| 72 | + self.assertEqual(config.debug_assert, original_value) | ||
| 73 | + | ||
| 74 | + def test_patch_with_tensor_device(self): | ||
| 75 | + """Test that config.patch works correctly with NPU tensors in the context""" | ||
| 76 | + x = torch.randn(3, 4).to(device_type) | ||
| 77 | + original_value = config.debug_assert | ||
| 78 | + with config.patch("debug_assert", not original_value): | ||
| 79 | + # Verify NPU tensor operations still work | ||
| 80 | + y = x.mm(x.T) | ||
| 81 | + self.assertEqual(y.device.type, device_type) | ||
| 82 | + self.assertEqual(config.debug_assert, not original_value) | ||
| 83 | + self.assertEqual(config.debug_assert, original_value) | ||
| 84 | + | ||
| 85 | + def test_patch_invalid_key(self): | ||
| 86 | + """Test patch with invalid key raises AttributeError.""" | ||
| 87 | + with self.assertRaises(AttributeError): | ||
| 88 | + with config.patch("non_existent_key", 42): | ||
| 89 | + pass | ||
| 90 | + | ||
| 91 | + def test_patch_invalid_dict_key(self): | ||
| 92 | + """Test patch dict with invalid key among valid keys raises AttributeError and does not change valid keys.""" | ||
| 93 | + original_value = config.debug_assert | ||
| 94 | + with self.assertRaises(AttributeError): | ||
| 95 | + with config.patch({"debug_assert": True, "invalid_key": 42}): | ||
| 96 | + pass | ||
| 97 | + # Verify valid key is unchanged after the failed patch | ||
| 98 | + self.assertEqual(config.debug_assert, original_value) | ||
| 99 | + | ||
| 100 | + | ||
| 101 | +if __name__ == "__main__": | ||
| 102 | + run_tests() | ||
| @@ -0,0 +1,184 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved. | ||
| 2 | +# | ||
| 3 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 4 | +# you may not use this file except in compliance with the License. | ||
| 5 | +# You may obtain a copy of the License at | ||
| 6 | +# | ||
| 7 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 8 | +# | ||
| 9 | +# Unless required by applicable law or agreed to in writing, software | ||
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +# See the License for the specific language governing permissions and | ||
| 13 | +# limitations under the License. | ||
| 14 | + | ||
| 15 | +""" | ||
| 16 | +Add validation cases for torch._functorch.vmap._add_batch_dim API on NPU: | ||
| 17 | +1. PyTorch community lacks sufficient and direct API validations for this API, so this file is added. | ||
| 18 | +2. This file validates torch._functorch.vmap._add_batch_dim (extendable). | ||
| 19 | +""" | ||
| 20 | +import torch | ||
| 21 | +from torch._functorch.vmap import _add_batch_dim | ||
| 22 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 23 | + | ||
| 24 | +device_type = ( | ||
| 25 | + acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 26 | +) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +class TestVmapAddBatchDim(TestCase): | ||
| 30 | + def test_add_batch_dim_basic(self): | ||
| 31 | + """Test basic _add_batch_dim functionality on NPU""" | ||
| 32 | + x = torch.randn(3, 4).to(device_type) | ||
| 33 | + vmap_level = 0 | ||
| 34 | + batched = _add_batch_dim(x, 0, vmap_level) | ||
| 35 | + self.assertIsNotNone(batched) | ||
| 36 | + self.assertIsInstance(batched, torch.Tensor) | ||
| 37 | + self.assertEqual(batched.device.type, device_type) | ||
| 38 | + self.assertEqual(batched.shape, (4,)) | ||
| 39 | + | ||
| 40 | + def test_add_batch_dim_with_vmap(self): | ||
| 41 | + """Test _add_batch_dim works correctly with vmap on NPU""" | ||
| 42 | + x = torch.randn(2, 3).to(device_type) | ||
| 43 | + y = torch.randn(2, 3).to(device_type) | ||
| 44 | + | ||
| 45 | + def dot_row(a, b): | ||
| 46 | + return (a * b).sum(dim=-1) | ||
| 47 | + | ||
| 48 | + result = torch.vmap(dot_row)(x, y) | ||
| 49 | + expected = dot_row(x, y) | ||
| 50 | + self.assertEqual(result, expected) | ||
| 51 | + self.assertEqual(result.device.type, device_type) | ||
| 52 | + | ||
| 53 | + def test_add_batch_dim_nested_vmap(self): | ||
| 54 | + """Test nested vmap with _add_batch_dim on NPU""" | ||
| 55 | + x = torch.randn(2, 3, 4).to(device_type) | ||
| 56 | + y = torch.randn(2, 3, 4).to(device_type) | ||
| 57 | + | ||
| 58 | + def matmul_row(a, b): | ||
| 59 | + return (a * b).sum(dim=-1) | ||
| 60 | + | ||
| 61 | + result = torch.vmap(torch.vmap(matmul_row))(x, y) | ||
| 62 | + self.assertEqual(result.shape, (2, 3)) | ||
| 63 | + self.assertEqual(result.device.type, device_type) | ||
| 64 | + | ||
| 65 | + def test_add_batch_dim_with_model(self): | ||
| 66 | + """Test _add_batch_dim with a simple model on NPU""" | ||
| 67 | + model = torch.nn.Linear(4, 2).to(device_type) | ||
| 68 | + x = torch.randn(3, 4).to(device_type) | ||
| 69 | + | ||
| 70 | + result = torch.vmap(lambda x: model(x))(x) | ||
| 71 | + expected = model(x) | ||
| 72 | + self.assertEqual(result, expected) | ||
| 73 | + self.assertEqual(result.device.type, device_type) | ||
| 74 | + | ||
| 75 | + def test_add_batch_dim_in_dims(self): | ||
| 76 | + """Test _add_batch_dim with different in_dims on NPU""" | ||
| 77 | + x = torch.randn(3, 4, 5).to(device_type) | ||
| 78 | + | ||
| 79 | + def identity(x): | ||
| 80 | + return x | ||
| 81 | + | ||
| 82 | + # Test in_dims=0 (default) | ||
| 83 | + result0 = torch.vmap(identity, in_dims=0)(x) | ||
| 84 | + self.assertEqual(result0.shape, (3, 4, 5)) | ||
| 85 | + | ||
| 86 | + # Test in_dims=1 | ||
| 87 | + result1 = torch.vmap(identity, in_dims=1)(x) | ||
| 88 | + self.assertEqual(result1.shape, (4, 3, 5)) | ||
| 89 | + | ||
| 90 | + # Test in_dims=-1 | ||
| 91 | + result_neg1 = torch.vmap(identity, in_dims=-1)(x) | ||
| 92 | + self.assertEqual(result_neg1.shape, (5, 3, 4)) | ||
| 93 | + | ||
| 94 | + def test_add_batch_dim_out_dims(self): | ||
| 95 | + """Test _add_batch_dim with different out_dims on NPU""" | ||
| 96 | + x = torch.randn(3, 4).to(device_type) | ||
| 97 | + | ||
| 98 | + def identity(x): | ||
| 99 | + return x | ||
| 100 | + | ||
| 101 | + # Test out_dims=0 (default) | ||
| 102 | + result0 = torch.vmap(identity, out_dims=0)(x) | ||
| 103 | + self.assertEqual(result0.shape, (3, 4)) | ||
| 104 | + | ||
| 105 | + # Test out_dims=1 | ||
| 106 | + result1 = torch.vmap(identity, out_dims=1)(x) | ||
| 107 | + self.assertEqual(result1.shape, (4, 3)) | ||
| 108 | + | ||
| 109 | + def test_add_batch_dim_with_grad(self): | ||
| 110 | + """Test _add_batch_dim works with gradient computation on NPU""" | ||
| 111 | + x = torch.randn(3, 3, device=device_type, requires_grad=True) | ||
| 112 | + w = torch.randn(3, 3, device=device_type, requires_grad=True) | ||
| 113 | + | ||
| 114 | + def fn(x, w): | ||
| 115 | + return (x * w).sum(dim=-1) | ||
| 116 | + | ||
| 117 | + result = torch.vmap(fn)(x, w) | ||
| 118 | + loss = result.sum() | ||
| 119 | + loss.backward() | ||
| 120 | + self.assertIsNotNone(x.grad) | ||
| 121 | + self.assertIsNotNone(w.grad) | ||
| 122 | + self.assertEqual(x.grad.device.type, device_type) | ||
| 123 | + self.assertEqual(w.grad.device.type, device_type) | ||
| 124 | + | ||
| 125 | + def test_add_batch_dim_direct_3d_batch_dim_0(self): | ||
| 126 | + """Direct API call: 3D tensor with batch_dim=0 returns shape=(4,5) (inner value sliced at dim 0).""" | ||
| 127 | + x = torch.randn(3, 4, 5, dtype=torch.float32).to(device_type) | ||
| 128 | + vmap_level = 0 | ||
| 129 | + batched = _add_batch_dim(x, 0, vmap_level) | ||
| 130 | + self.assertIsNotNone(batched) | ||
| 131 | + self.assertEqual(batched.shape, (4, 5)) | ||
| 132 | + self.assertEqual(batched.device.type, device_type) | ||
| 133 | + self.assertEqual(batched.dtype, torch.float32) | ||
| 134 | + | ||
| 135 | + def test_add_batch_dim_direct_3d_batch_dim_1(self): | ||
| 136 | + """Direct API call: 3D tensor with batch_dim=1 returns shape=(3,5).""" | ||
| 137 | + x = torch.randn(3, 4, 5, dtype=torch.float32).to(device_type) | ||
| 138 | + vmap_level = 0 | ||
| 139 | + batched = _add_batch_dim(x, 1, vmap_level) | ||
| 140 | + self.assertIsNotNone(batched) | ||
| 141 | + self.assertEqual(batched.shape, (3, 5)) | ||
| 142 | + self.assertEqual(batched.device.type, device_type) | ||
| 143 | + | ||
| 144 | + def test_add_batch_dim_direct_3d_batch_dim_2(self): | ||
| 145 | + """Direct API call: 3D tensor with batch_dim=2 returns shape=(3,4).""" | ||
| 146 | + x = torch.randn(3, 4, 5, dtype=torch.float32).to(device_type) | ||
| 147 | + vmap_level = 0 | ||
| 148 | + batched = _add_batch_dim(x, 2, vmap_level) | ||
| 149 | + self.assertIsNotNone(batched) | ||
| 150 | + self.assertEqual(batched.shape, (3, 4)) | ||
| 151 | + self.assertEqual(batched.device.type, device_type) | ||
| 152 | + | ||
| 153 | + def test_add_batch_dim_direct_preserves_dtype_and_device(self): | ||
| 154 | + """Direct API call preserves dtype and device.""" | ||
| 155 | + x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32).to(device_type) | ||
| 156 | + vmap_level = 0 | ||
| 157 | + batched = _add_batch_dim(x, 0, vmap_level) | ||
| 158 | + self.assertEqual(batched.dtype, torch.float32) | ||
| 159 | + self.assertEqual(batched.device.type, device_type) | ||
| 160 | + self.assertEqual(batched.shape, (2,)) | ||
| 161 | + | ||
| 162 | + def test_add_batch_dim_invalid_batch_dim(self): | ||
| 163 | + """Direct API call with out-of-range batch_dim wraps (x.ndim % batch_dim).""" | ||
| 164 | + x = torch.randn(3, 4).to(device_type) | ||
| 165 | + # batch_dim=5 on 2D tensor acts as batch_dim=1 (5 % 2 ≡ 1) | ||
| 166 | + batched = _add_batch_dim(x, 5, 0) | ||
| 167 | + self.assertEqual(batched.shape, (3,)) | ||
| 168 | + self.assertEqual(batched.device.type, device_type) | ||
| 169 | + | ||
| 170 | + def test_add_batch_dim_multiple_levels(self): | ||
| 171 | + """Test different vmap_levels produce independent batch dimensions.""" | ||
| 172 | + from torch._functorch.vmap import _remove_batch_dim | ||
| 173 | + x = torch.randn(3, 4).to(device_type) | ||
| 174 | + # Level 0 | ||
| 175 | + b_l0 = _add_batch_dim(x, 0, 0) | ||
| 176 | + # Level 1 - independent from level 0 | ||
| 177 | + b_l1 = _add_batch_dim(x, 0, 1) | ||
| 178 | + # Both levels show the same shape on the surface | ||
| 179 | + self.assertEqual(b_l0.shape, (4,)) | ||
| 180 | + self.assertEqual(b_l1.device.type, device_type) | ||
| 181 | + | ||
| 182 | + | ||
| 183 | +if __name__ == "__main__": | ||
| 184 | + run_tests() | ||