| 343 | |
| 344 | @eval_modes() |
| 345 | def test_batch_partial_slice(): |
| 346 | b = ndd.as_batch( |
| 347 | [ |
| 348 | ndd.tensor([[1, 2, 3], [4, 5, 6]], dtype=ndd.int32), |
| 349 | ndd.tensor([[7, 8, 9], [10, 11, 12]], dtype=ndd.int32), |
| 350 | ], |
| 351 | layout="XY", |
| 352 | ) |
| 353 | b11 = b.slice[..., 1] |
| 354 | assert b11.layout == "X" |
| 355 | assert b11.dtype == ndd.int32 |
| 356 | assert isinstance(b11, ndd.Batch) |
| 357 | assert np.array_equal(asnumpy(b11.tensors[0]), np.array([2, 5], dtype=np.int32)) |
| 358 | assert np.array_equal(asnumpy(b11.tensors[1]), np.array([8, 11], dtype=np.int32)) |
| 359 | |
| 360 | |
| 361 | @eval_modes() |