| 331 | } |
| 332 | template <class Src, class Dst> |
| 333 | void run_negative_index_test() { |
| 334 | constexpr size_t S0 = 200, S1 = 200; |
| 335 | HostTensorND hv_init{CompNode::load("xpu0"), dtype::Float32()}; |
| 336 | hv_init.resize({S0, S1}); |
| 337 | for (size_t i = 0; i < S0 * S1; ++i) |
| 338 | hv_init.ptr<float>()[i] = i; |
| 339 | |
| 340 | Src src; |
| 341 | Src src_sub; |
| 342 | Dst dst; |
| 343 | auto check = [&](size_t begin, size_t end, int axis) { |
| 344 | src.copy_from(hv_init).sync(); |
| 345 | src_sub = src.sub(Slice(begin, end).apply(src.layout(), axis)); |
| 346 | dst.copy_from(src_sub).sync(); |
| 347 | if (axis < 0) |
| 348 | axis += 2; |
| 349 | ASSERT_EQ(dst.layout().ndim, 2u); |
| 350 | for (int i = 0; i < 2; i++) { |
| 351 | if (i == axis) |
| 352 | ASSERT_EQ(dst.layout()[i], end - begin); |
| 353 | else |
| 354 | ASSERT_EQ(dst.layout()[i], 200u); |
| 355 | } |
| 356 | }; |
| 357 | check(100, 200, -1); |
| 358 | check(10, 20, -1); |
| 359 | check(100, 200, -2); |
| 360 | check(10, 20, -2); |
| 361 | EXPECT_THROW(check(100, 200, -3), MegBrainError); |
| 362 | EXPECT_THROW(check(10, 20, -3), MegBrainError); |
| 363 | EXPECT_THROW(check(100, 200, 2), MegBrainError); |
| 364 | EXPECT_THROW(check(10, 20, 2), MegBrainError); |
| 365 | } |
| 366 | |
| 367 | TEST(TestTensor, NegativeIndex) { |
| 368 | run_negative_index_test<HostTensorND, HostTensorND>(); |