已合并
test: add unflatten_dense_tensors NPU coverage #41632
yulin520创建于 7月14日
test: add unflatten_dense_tensors NPU coverage #41632
已合并
共 1 个文件变更+211-0
| @@ -0,0 +1,211 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or | ||
| 13 | +# implied. | ||
| 14 | +# See the License for the specific language governing permissions and | ||
| 15 | +# limitations under the License. | ||
| 16 | + | ||
| 17 | +""" | ||
| 18 | +Add validation cases for torch.nn APIs on NPU: | ||
| 19 | +1. PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. | ||
| 20 | +2. This file validates torch._C._nn.unflatten_dense_tensors (extendable). | ||
| 21 | +""" | ||
| 22 | + | ||
| 23 | +import torch | ||
| 24 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 25 | + | ||
| 26 | +device_type = ( | ||
| 27 | + acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 28 | +) | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +class TestUnflattenDenseTensors(TestCase): | ||
群 | |||
| 32 | + def test_unflatten_dense_tensors_basic(self): | ||
| 33 | + flat = torch.arange( | ||
| 34 | + 12, | ||
| 35 | + dtype=torch.float32, | ||
| 36 | + device=device_type, | ||
| 37 | + ) | ||
| 38 | + | ||
| 39 | + tensors = [ | ||
| 40 | + torch.empty( | ||
| 41 | + (2, 3), | ||
| 42 | + dtype=torch.float32, | ||
| 43 | + device=device_type, | ||
| 44 | + ), | ||
| 45 | + torch.empty( | ||
| 46 | + (3, 2), | ||
| 47 | + dtype=torch.float32, | ||
| 48 | + device=device_type, | ||
| 49 | + ), | ||
| 50 | + ] | ||
| 51 | + | ||
| 52 | + outputs = torch._C._nn.unflatten_dense_tensors( | ||
| 53 | + flat, | ||
| 54 | + tensors, | ||
| 55 | + ) | ||
| 56 | + | ||
| 57 | + self.assertIsInstance(outputs, tuple) | ||
| 58 | + self.assertEqual(len(outputs), 2) | ||
| 59 | + | ||
| 60 | + self.assertEqual(tuple(outputs[0].shape), (2, 3)) | ||
| 61 | + self.assertEqual(tuple(outputs[1].shape), (3, 2)) | ||
| 62 | + | ||
| 63 | + self.assertEqual(outputs[0].device.type, device_type) | ||
| 64 | + self.assertEqual(outputs[1].device.type, device_type) | ||
| 65 | + | ||
| 66 | + self.assertEqual(outputs[0].dtype, torch.float32) | ||
| 67 | + self.assertEqual(outputs[1].dtype, torch.float32) | ||
| 68 | + | ||
| 69 | + expected_first = torch.arange( | ||
| 70 | + 0, | ||
| 71 | + 6, | ||
| 72 | + dtype=torch.float32, | ||
| 73 | + device=device_type, | ||
| 74 | + ).reshape(2, 3) | ||
| 75 | + | ||
| 76 | + expected_second = torch.arange( | ||
| 77 | + 6, | ||
| 78 | + 12, | ||
| 79 | + dtype=torch.float32, | ||
| 80 | + device=device_type, | ||
| 81 | + ).reshape(3, 2) | ||
| 82 | + | ||
| 83 | + self.assertTrue( | ||
| 84 | + torch.equal(outputs[0], expected_first) | ||
| 85 | + ) | ||
| 86 | + self.assertTrue( | ||
| 87 | + torch.equal(outputs[1], expected_second) | ||
| 88 | + ) | ||
| 89 | + | ||
| 90 | + flat_storage_ptr = ( | ||
| 91 | + flat.untyped_storage().data_ptr() | ||
| 92 | + ) | ||
| 93 | + | ||
| 94 | + self.assertEqual( | ||
| 95 | + outputs[0].untyped_storage().data_ptr(), | ||
| 96 | + flat_storage_ptr, | ||
| 97 | + ) | ||
| 98 | + self.assertEqual( | ||
| 99 | + outputs[1].untyped_storage().data_ptr(), | ||
| 100 | + flat_storage_ptr, | ||
| 101 | + ) | ||
| 102 | + | ||
| 103 | + self.assertEqual(outputs[0].storage_offset(), 0) | ||
| 104 | + self.assertEqual(outputs[1].storage_offset(), 6) | ||
| 105 | + | ||
| 106 | + def test_unflatten_dense_tensors_with_empty_tensor(self): | ||
| 107 | + flat = torch.arange( | ||
| 108 | + 6, | ||
| 109 | + dtype=torch.float32, | ||
| 110 | + device=device_type, | ||
| 111 | + ) | ||
| 112 | + | ||
| 113 | + tensors = [ | ||
| 114 | + torch.empty( | ||
| 115 | + (0,), | ||
| 116 | + dtype=torch.float32, | ||
| 117 | + device=device_type, | ||
| 118 | + ), | ||
| 119 | + torch.empty( | ||
| 120 | + (2, 3), | ||
| 121 | + dtype=torch.float32, | ||
| 122 | + device=device_type, | ||
| 123 | + ), | ||
| 124 | + ] | ||
| 125 | + | ||
| 126 | + outputs = torch._C._nn.unflatten_dense_tensors( | ||
| 127 | + flat, | ||
| 128 | + tensors, | ||
| 129 | + ) | ||
| 130 | + | ||
| 131 | + self.assertIsInstance(outputs, tuple) | ||
| 132 | + self.assertEqual(len(outputs), 2) | ||
| 133 | + | ||
| 134 | + self.assertEqual(tuple(outputs[0].shape), (0,)) | ||
| 135 | + self.assertEqual(outputs[0].numel(), 0) | ||
| 136 | + self.assertEqual(outputs[0].device.type, device_type) | ||
| 137 | + | ||
| 138 | + self.assertEqual(tuple(outputs[1].shape), (2, 3)) | ||
| 139 | + self.assertEqual(outputs[1].device.type, device_type) | ||
| 140 | + | ||
| 141 | + expected = torch.arange( | ||
| 142 | + 6, | ||
| 143 | + dtype=torch.float32, | ||
| 144 | + device=device_type, | ||
| 145 | + ).reshape(2, 3) | ||
| 146 | + | ||
| 147 | + self.assertTrue( | ||
| 148 | + torch.equal(outputs[1], expected) | ||
| 149 | + ) | ||
| 150 | + | ||
| 151 | + self.assertEqual( | ||
| 152 | + outputs[1].untyped_storage().data_ptr(), | ||
| 153 | + flat.untyped_storage().data_ptr(), | ||
| 154 | + ) | ||
| 155 | + self.assertEqual(outputs[1].storage_offset(), 0) | ||
| 156 | + | ||
| 157 | + def test_unflatten_dense_tensors_autograd(self): | ||
| 158 | + flat = torch.arange( | ||
| 159 | + 10, | ||
| 160 | + dtype=torch.float32, | ||
| 161 | + device=device_type, | ||
| 162 | + requires_grad=True, | ||
| 163 | + ) | ||
| 164 | + | ||
| 165 | + tensors = [ | ||
| 166 | + torch.empty((2, 3), dtype=torch.float32, device=device_type), | ||
| 167 | + torch.empty((4,), dtype=torch.float32, device=device_type), | ||
| 168 | + ] | ||
| 169 | + | ||
| 170 | + outputs = torch._C._nn.unflatten_dense_tensors(flat, tensors) | ||
| 171 | + | ||
| 172 | + self.assertEqual(len(outputs), 2) | ||
| 173 | + self.assertTrue(outputs[0].requires_grad) | ||
| 174 | + self.assertTrue(outputs[1].requires_grad) | ||
| 175 | + self.assertEqual(outputs[0].device.type, device_type) | ||
| 176 | + self.assertEqual(outputs[1].device.type, device_type) | ||
| 177 | + | ||
| 178 | + loss = outputs[0].sum() + (outputs[1] * 2).sum() | ||
| 179 | + loss.backward() | ||
| 180 | + | ||
| 181 | + expected_grad = torch.tensor( | ||
| 182 | + [1.0] * 6 + [2.0] * 4, | ||
| 183 | + dtype=torch.float32, | ||
| 184 | + device=device_type, | ||
| 185 | + ) | ||
| 186 | + | ||
| 187 | + self.assertIsNotNone(flat.grad) | ||
| 188 | + self.assertEqual(flat.grad.device.type, device_type) | ||
| 189 | + self.assertEqual(flat.grad, expected_grad) | ||
| 190 | + | ||
| 191 | + def test_unflatten_dense_tensors_dtype_preservation(self): | ||
| 192 | + for dtype in (torch.float16, torch.float32): | ||
| 193 | + with self.subTest(dtype=dtype): | ||
| 194 | + flat = torch.tensor( | ||
| 195 | + [0, 1, 2, 3, 4, 5], | ||
| 196 | + dtype=dtype, | ||
| 197 | + device=device_type, | ||
| 198 | + ) | ||
| 199 | + tensors = [ | ||
| 200 | + torch.empty((2,), dtype=dtype, device=device_type), | ||
| 201 | + torch.empty((2, 2), dtype=dtype, device=device_type), | ||
| 202 | + ] | ||
| 203 | + | ||
| 204 | + outputs = torch._C._nn.unflatten_dense_tensors(flat, tensors) | ||
| 205 | + | ||
| 206 | + self.assertEqual(outputs[0].dtype, dtype) | ||
| 207 | + self.assertEqual(outputs[1].dtype, dtype) | ||
| 208 | + | ||
| 209 | + | ||
| 210 | +if __name__ == "__main__": | ||
| 211 | + run_tests() | ||
用例只覆盖 float32，缺少对 dtype 保持契约的基本验证