已合并
test(npu): add pytree supported nodes pop coverage #41812
jerry lee创建于 7月16日
test(npu): add pytree supported nodes pop coverage #41812
已合并
共 1 个文件变更+47-0
| @@ -0,0 +1,47 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch.utils._pytree APIs on NPU: | ||
| 3 | +1. PyTorch community lacks sufficient and direct validation for torch.utils._pytree.SUPPORTED_NODES.pop. | ||
| 4 | +2. This file validates pytree registry mutation behavior and is extendable. | ||
| 5 | +""" | ||
| 6 | + | ||
| 7 | +import torch | ||
| 8 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 9 | +from torch.utils import _pytree | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +class TestPytreeSupportedNodes(TestCase): | ||
| 16 | + def test_pop(self): | ||
| 17 | + class CustomNode: | ||
| 18 | + def __init__(self, value): | ||
| 19 | + self.value = value | ||
| 20 | + | ||
| 21 | + def flatten_fn(node): | ||
| 22 | + return [node.value], None | ||
| 23 | + | ||
| 24 | + def unflatten_fn(values, context): | ||
| 25 | + return CustomNode(values[0]) | ||
| 26 | + | ||
| 27 | + _pytree.register_pytree_node(CustomNode, flatten_fn, unflatten_fn) | ||
| 28 | + self.addCleanup(_pytree._deregister_pytree_node, CustomNode) | ||
| 29 | + value = torch.tensor([1.0]).to(device_type) | ||
| 30 | + node = CustomNode(value) | ||
| 31 | + | ||
| 32 | + leaves, _ = _pytree.tree_flatten(node) | ||
| 33 | + self.assertEqual(leaves, [value]) | ||
| 34 | + | ||
| 35 | + node_def = _pytree.SUPPORTED_NODES.pop(CustomNode) | ||
| 36 | + self.addCleanup(_pytree.SUPPORTED_NODES.__setitem__, CustomNode, node_def) | ||
| 37 | + self.assertIs(node_def.flatten_fn, flatten_fn) | ||
| 38 | + self.assertNotIn(CustomNode, _pytree.SUPPORTED_NODES) | ||
| 39 | + self.assertIsNone(_pytree.SUPPORTED_NODES.pop(CustomNode, None)) | ||
| 40 | + | ||
| 41 | + leaves, spec = _pytree.tree_flatten(node) | ||
| 42 | + self.assertTrue(spec.is_leaf()) | ||
| 43 | + self.assertEqual(leaves, [node]) | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +if __name__ == "__main__": | ||
| 47 | + run_tests() | ||