| 360 | |
| 361 | @eval_modes() |
| 362 | def test_batch_slice(): |
| 363 | b = ndd.as_batch( |
| 364 | [ |
| 365 | ndd.tensor([[1, 2, 3], [4, 5, 6]], dtype=ndd.uint16), |
| 366 | ndd.tensor([[7, 8, 9, 10], [11, 12, 13, 14]], dtype=ndd.uint16), |
| 367 | ], |
| 368 | layout="XY", |
| 369 | ) |
| 370 | sliced = b.slice[..., 1:-1] |
| 371 | assert sliced.layout == "XY" |
| 372 | assert sliced.dtype == ndd.uint16 |
| 373 | assert np.array_equal(asnumpy(sliced.tensors[0]), np.array([[2], [5]], dtype=np.uint16)) |
| 374 | assert np.array_equal(asnumpy(sliced.tensors[1]), np.array([[8, 9], [12, 13]], dtype=np.uint16)) |
| 375 | |
| 376 | |
| 377 | @eval_modes() |