已合并
【API一致性任务】test: add torch._C._functorch.is_batchedtensor validation cases on NPU #43533
cuiyunhao-2026创建于 8月1日
【API一致性任务】test: add torch._C._functorch.is_batchedtensor validation cases on NPU #43533
已合并
共 1 个文件变更+133-0
| @@ -0,0 +1,133 @@ | |||
| 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 implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +""" | ||
| 17 | +Add validation cases for torch._C._functorch APIs on NPU: | ||
| 18 | +1. PyTorch community lacks direct test cases for the following APIs, | ||
| 19 | + so this file is added. | ||
| 20 | +2. This file validates torch._C._functorch.is_batchedtensor (extendable). | ||
| 21 | +""" | ||
| 22 | + | ||
| 23 | +import torch | ||
| 24 | +from torch._C._functorch import ( | ||
| 25 | + _add_batch_dim, | ||
| 26 | + _vmap_decrement_nesting, | ||
| 27 | + _vmap_increment_nesting, | ||
| 28 | + get_unwrapped, | ||
| 29 | + is_batchedtensor, | ||
| 30 | +) | ||
| 31 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 35 | + | ||
| 36 | + | ||
| 37 | +class TestFunctorchIsBatchedTensor(TestCase): | ||
| 38 | + | ||
| 39 | + def test_is_batchedtensor_plain_tensor(self): | ||
| 40 | + # a normal tensor is not a BatchedTensor | ||
| 41 | + x = torch.randn(2, 3).to(device_type) | ||
| 42 | + self.assertFalse(is_batchedtensor(x)) | ||
| 43 | + | ||
| 44 | + def test_is_batchedtensor_inside_vmap(self): | ||
| 45 | + # the tensor passed into a vmap-ed function is a BatchedTensor | ||
| 46 | + seen = [] | ||
| 47 | + | ||
| 48 | + def fn(t): | ||
| 49 | + seen.append(is_batchedtensor(t)) | ||
| 50 | + return t.sum() | ||
| 51 | + | ||
| 52 | + x = torch.randn(4, 3).to(device_type) | ||
| 53 | + torch.vmap(fn)(x) | ||
| 54 | + self.assertEqual(seen, [True]) | ||
| 55 | + | ||
| 56 | + def test_is_batchedtensor_manual_batch_dim(self): | ||
| 57 | + # manually wrapped BatchedTensor is recognized, unwrapping restores False | ||
| 58 | + x = torch.randn(3, 5).to(device_type) | ||
| 59 | + level = _vmap_increment_nesting(3, "error") | ||
| 60 | + try: | ||
| 61 | + batched = _add_batch_dim(x, 0, level) | ||
| 62 | + self.assertTrue(is_batchedtensor(batched)) | ||
| 63 | + self.assertFalse(is_batchedtensor(get_unwrapped(batched))) | ||
| 64 | + finally: | ||
| 65 | + _vmap_decrement_nesting() | ||
| 66 | + # The finally block must restore the vmap nesting level so later tests | ||
| 67 | + # are not polluted. Re-incrementing yields the same level we started | ||
| 68 | + # from, proving the nesting counter was fully restored. | ||
| 69 | + level_after = _vmap_increment_nesting(3, "error") | ||
| 70 | + try: | ||
| 71 | + self.assertEqual(level_after, level) | ||
| 72 | + finally: | ||
| 73 | + _vmap_decrement_nesting() | ||
| 74 | + | ||
| 75 | + def test_is_batchedtensor_nested_vmap(self): | ||
| 76 | + # BatchedTensor of nested vmap is still a BatchedTensor | ||
| 77 | + seen = [] | ||
| 78 | + | ||
| 79 | + def fn(t): | ||
| 80 | + seen.append(is_batchedtensor(t)) | ||
| 81 | + return t.sum() | ||
| 82 | + | ||
| 83 | + x = torch.randn(2, 3, 4).to(device_type) | ||
| 84 | + torch.vmap(torch.vmap(fn))(x) | ||
| 85 | + self.assertEqual(seen, [True]) | ||
| 86 | + | ||
| 87 | + def test_is_batchedtensor_outside_vmap(self): | ||
| 88 | + # the tensor is no longer batched after vmap returns | ||
| 89 | + x = torch.randn(4, 3).to(device_type) | ||
| 90 | + out = torch.vmap(lambda t: t * 2)(x) | ||
| 91 | + self.assertFalse(is_batchedtensor(out)) | ||
| 92 | + | ||
| 93 | + def test_is_batchedtensor_various_dtypes(self): | ||
| 94 | + # the result only depends on batching, not on dtype | ||
| 95 | + for dtype in (torch.float32, torch.float16, torch.int32, torch.bool): | ||
| 96 | + x = torch.ones(2, 3, dtype=dtype).to(device_type) | ||
| 97 | + self.assertFalse(is_batchedtensor(x)) | ||
| 98 | + | ||
| 99 | + def test_is_batchedtensor_non_tensor_input(self): | ||
| 100 | + # non-tensor input is rejected | ||
| 101 | + with self.assertRaises(TypeError): | ||
| 102 | + is_batchedtensor(1) | ||
| 103 | + | ||
| 104 | + def test_is_batchedtensor_zero_dim_tensor(self): | ||
| 105 | + # a 0-d (scalar) tensor is not a BatchedTensor | ||
| 106 | + x = torch.tensor(7).to(device_type) | ||
| 107 | + self.assertFalse(is_batchedtensor(x)) | ||
| 108 | + | ||
| 109 | + def test_is_batchedtensor_illegal_inputs(self): | ||
| 110 | + # None / str / arbitrary object inputs are rejected | ||
| 111 | + for bad in (None, "x", object()): | ||
| 112 | + with self.assertRaises(TypeError): | ||
| 113 | + is_batchedtensor(bad) | ||
| 114 | + | ||
| 115 | + def test_is_batchedtensor_other_device(self): | ||
| 116 | + # device of the tensor does not affect the batching check | ||
| 117 | + x_npu = torch.randn(2, 3).to(device_type) | ||
| 118 | + x_cpu = torch.randn(2, 3) | ||
| 119 | + self.assertFalse(is_batchedtensor(x_npu)) | ||
| 120 | + self.assertFalse(is_batchedtensor(x_cpu)) | ||
| 121 | + | ||
| 122 | + def test_is_batchedtensor_nesting_balance(self): | ||
| 123 | + # A full increment/decrement cycle must restore the vmap nesting level, | ||
| 124 | + # guarding against corrupted nesting state. | ||
| 125 | + level = _vmap_increment_nesting(2, "error") | ||
| 126 | + try: | ||
| 127 | + self.assertGreaterEqual(level, 1) | ||
| 128 | + finally: | ||
| 129 | + _vmap_decrement_nesting() | ||
| 130 | + | ||
| 131 | + | ||
| 132 | +if __name__ == "__main__": | ||
| 133 | + run_tests() | ||