| 32 | @eval_modes() |
| 33 | @params(("cpu",), ("gpu",)) |
| 34 | def test_batch_construction(device_type): |
| 35 | t0 = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32) |
| 36 | t1 = np.array([[7, 8, 9], [10, 11, 12]], dtype=np.int32) |
| 37 | |
| 38 | b = ndd.batch( |
| 39 | [ |
| 40 | t0, |
| 41 | t1, |
| 42 | ], |
| 43 | device=ndd.Device(device_type), |
| 44 | layout="AB", |
| 45 | ) |
| 46 | |
| 47 | assert isinstance(b, ndd.Batch) |
| 48 | assert b.device.device_type == device_type |
| 49 | assert b.layout == "AB" |
| 50 | assert np.array_equal(asnumpy(b.tensors[0]), t0) |
| 51 | assert np.array_equal(asnumpy(b.tensors[1]), t1) |
| 52 | # check that modifying the original arrays doesn't affect the batch |
| 53 | t0[0, 0] += 1 |
| 54 | t1[0, 0] += 1 |
| 55 | assert not np.array_equal(asnumpy(b.tensors[0]), t0) |
| 56 | assert not np.array_equal(asnumpy(b.tensors[1]), t1) |
| 57 | |
| 58 | b.evaluate() |
| 59 | assert b._storage.layout() == "AB" |
| 60 | |
| 61 | |
| 62 | @eval_modes() |