"""
测试目的：验证 torch.chunk 接口功能正确性
API 名称：torch.chunk
API 签名：torch.chunk(input, chunks, dim=0) -> tuple[Tensor, ...]
覆盖维度表：
| 覆盖维度 | 说明 | 覆盖情况 |
|------------------|--------------------------------------------------------------|------------------------------------------------|
| 空/非空 | 沿拼接维 size 为 0 的输入 | 已覆盖 |
| 枚举选项 | dim 取 0、正索引、负索引；chunks 取 1、>1 | 已覆盖 |
| 参数类型 | input 为 Tensor；chunks 为 int；dim 为 int | 已覆盖 |
| 传参与不传参 | dim 省略默认 0 | 已覆盖 |
| 等价类/边界值 | 可整除与不可整除的切分、高维、非连续输入 | 已覆盖 |
| 正常传参场景 | NPU 上典型 shape / dtype，返回 tuple 且子张量 device/dtype 一致 | 已覆盖 |
| 异常传参场景 | chunks<=0、非法 dim | 已覆盖 |
| 混合设备输入 | 单 Tensor 输入，不适用 | 不适用 |
未覆盖项及原因：
- 无
注意：本测试仅验证功能正确性（调用不报错、输出 shape/dtype/device/类型符合预期），
不做精度和数值正确性校验。
"""
import torch
import torch_npu
try:
from torch_npu.testing.testcase import TestCase, run_tests
except ImportError:
import sys
import unittest
from unittest import TestCase
def run_tests():
unittest.main(argv=sys.argv)
class TestChunk(TestCase):
"""Functional tests for torch.chunk on NPU."""
def setUp(self):
super().setUp()
self.device_name = torch._C._get_privateuse1_backend_name()
self.assertEqual(
self.device_name,
"npu",
f"Expected device 'npu', got '{self.device_name}'",
)
self.device = torch.device(self.device_name)
def test_chunk_npu_dim0_equal_parts(self):
x = torch.randn(6, 4, device=self.device)
parts = torch.chunk(x, 3, dim=0)
self.assertEqual(len(parts), 3)
for p in parts:
self.assertEqual(p.shape, torch.Size([2, 4]))
self.assertEqual(p.dtype, torch.float32)
self.assertEqual(p.device.type, self.device_name)
def test_chunk_npu_dim1(self):
x = torch.randn(2, 6, device=self.device)
parts = torch.chunk(x, 2, dim=1)
self.assertEqual(len(parts), 2)
self.assertEqual(parts[0].shape, torch.Size([2, 3]))
self.assertEqual(parts[1].shape, torch.Size([2, 3]))
def test_chunk_npu_dim_negative(self):
x = torch.randn(2, 3, 8, device=self.device)
parts = torch.chunk(x, 2, dim=-1)
self.assertEqual(len(parts), 2)
self.assertEqual(parts[0].shape, torch.Size([2, 3, 4]))
self.assertEqual(parts[1].shape, torch.Size([2, 3, 4]))
def test_chunk_npu_chunks_one(self):
x = torch.randn(4, 5, device=self.device)
parts = torch.chunk(x, 1, dim=0)
self.assertEqual(len(parts), 1)
self.assertEqual(parts[0].shape, x.shape)
self.assertEqual(parts[0].dtype, x.dtype)
self.assertEqual(parts[0].device.type, x.device.type)
def test_chunk_npu_default_dim(self):
x = torch.randn(4, 3, device=self.device)
parts_default = torch.chunk(x, 2)
parts_explicit = torch.chunk(x, 2, dim=0)
self.assertEqual(len(parts_default), len(parts_explicit))
for a, b in zip(parts_default, parts_explicit):
self.assertEqual(a.shape, b.shape)
def test_chunk_npu_high_rank(self):
x = torch.randn(2, 3, 4, 5, 6, device=self.device)
parts = torch.chunk(x, 2, dim=2)
self.assertEqual(len(parts), 2)
self.assertEqual(parts[0].shape, torch.Size([2, 3, 2, 5, 6]))
def test_chunk_npu_non_contiguous(self):
x = torch.randn(6, 4, device=self.device).t()
self.assertFalse(x.is_contiguous())
parts = torch.chunk(x, 3, dim=0)
total = 0
for p in parts:
total += p.shape[0]
self.assertEqual(p.device.type, self.device_name)
self.assertEqual(p.dtype, x.dtype)
self.assertEqual(total, x.shape[0])
out = torch.cat(list(parts), dim=0)
self.assertEqual(out.shape, x.shape)
def test_chunk_npu_uneven_split(self):
x = torch.randn(5, 2, device=self.device)
parts = torch.chunk(x, 2, dim=0)
self.assertEqual(len(parts), 2)
self.assertEqual(parts[0].shape[0] + parts[1].shape[0], 5)
self.assertEqual(parts[0].shape[1], 2)
self.assertEqual(parts[1].shape[1], 2)
def test_chunk_npu_empty_along_dim(self):
x = torch.empty(0, 3, device=self.device, dtype=torch.float32)
parts = torch.chunk(x, 2, dim=0)
self.assertEqual(len(parts), 2)
self.assertEqual(parts[0].shape[0] + parts[1].shape[0], 0)
def test_chunk_npu_supported_dtypes(self):
dtypes = [
torch.float32,
torch.float16,
torch.bfloat16,
torch.int32,
torch.int64,
torch.bool,
]
for dtype in dtypes:
if dtype == torch.bool:
x = torch.tensor([[True, False], [False, True]], device=self.device)
elif dtype in (torch.int32, torch.int64):
x = torch.tensor([[1, 2], [3, 4]], dtype=dtype, device=self.device)
else:
x = torch.ones(4, 2, dtype=dtype, device=self.device)
parts = torch.chunk(x, 2, dim=0)
self.assertEqual(len(parts), 2)
for p in parts:
self.assertEqual(p.dtype, dtype, f"dtype mismatch for {dtype}")
def test_chunk_npu_invalid_chunks_zero_raises(self):
x = torch.randn(2, 2, device=self.device)
with self.assertRaises(RuntimeError):
torch.chunk(x, 0, dim=0)
torch.npu.synchronize()
def test_chunk_npu_invalid_chunks_negative_raises(self):
x = torch.randn(2, 2, device=self.device)
with self.assertRaises(RuntimeError):
torch.chunk(x, -1, dim=0)
torch.npu.synchronize()
def test_chunk_npu_invalid_dim_raises(self):
x = torch.randn(2, 3, device=self.device)
with self.assertRaises((IndexError, RuntimeError)):
torch.chunk(x, 2, dim=3)
torch.npu.synchronize()
def test_chunk_cpu_baseline(self):
x = torch.randn(6, 4)
parts = torch.chunk(x, 3, dim=0)
self.assertEqual(len(parts), 3)
self.assertEqual(parts[0].shape, torch.Size([2, 4]))
def test_chunk_cpu_baseline_dim1(self):
x = torch.randn(2, 6)
parts = torch.chunk(x, 2, dim=1)
self.assertEqual(len(parts), 2)
self.assertEqual(parts[0].shape, torch.Size([2, 3]))
if __name__ == "__main__":
run_tests()
