已合并
test(refs): add test cases for torch._refs._maybe_broadcast #42122
木路折创建于 7月19日
test(refs): add test cases for torch._refs._maybe_broadcast #42122
已合并
共 1 个文件变更+120-0
| @@ -0,0 +1,120 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 3 | +# All rights reserved. | ||
| 4 | +# | ||
| 5 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 6 | +# you may not use this file except in compliance with the License. | ||
| 7 | +# You may obtain a copy of the License at | ||
| 8 | +# | ||
| 9 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 10 | +# | ||
| 11 | +# Unless required by applicable law or agreed to in writing, software | ||
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 14 | +# See the License for the specific language governing permissions and | ||
| 15 | +# limitations under the License. | ||
| 16 | +# Owner(s): ["module: library"] | ||
| 17 | + | ||
| 18 | +"""Add validation cases for torch._refs._maybe_broadcast on NPU: | ||
| 19 | +1. PyTorch community lacks sufficient and direct API validations for | ||
| 20 | + torch._refs._maybe_broadcast, so this file is added. | ||
| 21 | +2. This file validates the broadcasting behavior of _maybe_broadcast for | ||
| 22 | + TensorLike inputs, including same-shape, compatible-shape, CPU-scalar | ||
| 23 | + preservation/expansion, and error cases (extendable). | ||
| 24 | +""" | ||
| 25 | + | ||
| 26 | +import torch | ||
| 27 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 28 | + | ||
| 29 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +class TestLibraryRefs(TestCase): | ||
| 33 | + """Test torch._refs._maybe_broadcast on NPU and CPU scalar tensors.""" | ||
| 34 | + | ||
| 35 | + def test_maybe_broadcast_same_shape_npu(self): | ||
| 36 | + """Tensors with the same NPU shape are returned unchanged.""" | ||
| 37 | + x = torch.empty(2, 3).to(device_type) | ||
| 38 | + y = torch.empty(2, 3).to(device_type) | ||
| 39 | + bx, by = torch._refs._maybe_broadcast(x, y) | ||
| 40 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 41 | + self.assertEqual(by.shape, (2, 3)) | ||
| 42 | + self.assertEqual(bx.device, x.device) | ||
| 43 | + self.assertEqual(by.device, y.device) | ||
| 44 | + self.assertIs(bx, x) | ||
| 45 | + self.assertIs(by, y) | ||
| 46 | + | ||
| 47 | + def test_maybe_broadcast_different_shapes_npu(self): | ||
| 48 | + """NPU tensors with compatible shapes are expanded to the common shape.""" | ||
| 49 | + x = torch.empty(1, 3).to(device_type) | ||
| 50 | + y = torch.empty(2, 1).to(device_type) | ||
| 51 | + bx, by = torch._refs._maybe_broadcast(x, y) | ||
| 52 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 53 | + self.assertEqual(by.shape, (2, 3)) | ||
| 54 | + self.assertEqual(bx.device, x.device) | ||
| 55 | + self.assertEqual(by.device, y.device) | ||
| 56 | + | ||
| 57 | + def test_maybe_broadcast_number(self): | ||
| 58 | + """A Number argument is returned as-is.""" | ||
| 59 | + x = torch.empty(2, 3).to(device_type) | ||
| 60 | + number = 2.0 | ||
| 61 | + bx, bnumber = torch._refs._maybe_broadcast(x, number) | ||
| 62 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 63 | + self.assertEqual(bx.device, x.device) | ||
| 64 | + self.assertEqual(bnumber, number) | ||
| 65 | + | ||
| 66 | + def test_maybe_broadcast_none(self): | ||
| 67 | + """A None argument is returned as None.""" | ||
| 68 | + x = torch.empty(2, 3).to(device_type) | ||
| 69 | + bx, bnone = torch._refs._maybe_broadcast(x, None) | ||
| 70 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 71 | + self.assertEqual(bx.device, x.device) | ||
| 72 | + self.assertIsNone(bnone) | ||
| 73 | + | ||
| 74 | + def test_maybe_broadcast_cpu_scalar_preserved(self): | ||
| 75 | + """CPU scalar tensors are preserved when preserve_cpu_scalar_tensors is True.""" | ||
| 76 | + x = torch.empty(2, 3).to(device_type) | ||
| 77 | + scalar = torch.tensor(1.0) # CPU scalar, intentionally not on NPU | ||
| 78 | + bx, bscalar = torch._refs._maybe_broadcast( | ||
| 79 | + x, scalar, preserve_cpu_scalar_tensors=True | ||
| 80 | + ) | ||
| 81 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 82 | + self.assertEqual(bx.device, x.device) | ||
| 83 | + self.assertEqual(bscalar.shape, ()) | ||
| 84 | + self.assertEqual(bscalar.device, torch.device("cpu")) | ||
| 85 | + self.assertIs(bscalar, scalar) | ||
| 86 | + | ||
| 87 | + def test_maybe_broadcast_cpu_scalar_expanded(self): | ||
| 88 | + """CPU scalar tensors are expanded when preserve_cpu_scalar_tensors is False.""" | ||
| 89 | + x = torch.empty(2, 3).to(device_type) | ||
| 90 | + scalar = torch.tensor(1.0) # CPU scalar, intentionally not on NPU | ||
| 91 | + bx, bscalar = torch._refs._maybe_broadcast( | ||
| 92 | + x, scalar, preserve_cpu_scalar_tensors=False | ||
| 93 | + ) | ||
| 94 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 95 | + self.assertEqual(bx.device, x.device) | ||
| 96 | + self.assertEqual(bscalar.shape, (2, 3)) | ||
| 97 | + self.assertEqual(bscalar.device, torch.device("cpu")) | ||
| 98 | + | ||
| 99 | + def test_maybe_broadcast_scalar_expanded(self): | ||
| 100 | + """A 0-dim scalar tensor is broadcast when CPU preservation is disabled.""" | ||
| 101 | + x = torch.empty(2, 3).to(device_type) | ||
| 102 | + scalar = torch.tensor(1.0).to(device_type) | ||
| 103 | + bx, bscalar = torch._refs._maybe_broadcast( | ||
| 104 | + x, scalar, preserve_cpu_scalar_tensors=False | ||
| 105 | + ) | ||
| 106 | + self.assertEqual(bx.shape, (2, 3)) | ||
| 107 | + self.assertEqual(bx.device, x.device) | ||
| 108 | + self.assertEqual(bscalar.shape, (2, 3)) | ||
| 109 | + self.assertEqual(bscalar.device, x.device) | ||
| 110 | + | ||
| 111 | + def test_maybe_broadcast_incompatible_shapes(self): | ||
| 112 | + """Incompatible tensor shapes raise RuntimeError.""" | ||
| 113 | + x = torch.empty(2, 3).to(device_type) | ||
| 114 | + y = torch.empty(3, 2).to(device_type) | ||
| 115 | + with self.assertRaises(RuntimeError): | ||
| 116 | + torch._refs._maybe_broadcast(x, y) | ||
| 117 | + | ||
| 118 | + | ||
| 119 | +if __name__ == "__main__": | ||
| 120 | + run_tests() | ||