(rank, axes, device, in_layout)
| 567 | |
| 568 | @nottest |
| 569 | def _test_reduce_large_data(rank, axes, device, in_layout): |
| 570 | batch_size = 16 |
| 571 | num_batches = 2 |
| 572 | data = fast_large_random_batches(rank, batch_size, num_batches) |
| 573 | |
| 574 | pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0 if device == "gpu" else None) |
| 575 | input = fn.external_source(data, cycle=True, device=device, layout=in_layout) |
| 576 | reduced = fn.reductions.sum(input, axes=axes) |
| 577 | pipe.set_outputs(reduced) |
| 578 | |
| 579 | for b, batch in enumerate(data): |
| 580 | (out,) = pipe.run() |
| 581 | check_layout(out, in_layout, axes, False) |
| 582 | if device == "gpu": |
| 583 | out = out.as_cpu() |
| 584 | for i in range(batch_size): |
| 585 | ref = np.sum(batch[i].astype(np.float64), axis=axes) |
| 586 | assert np.allclose(out[i], ref, 1e-5, 1e-5) |
| 587 | |
| 588 | |
| 589 | def test_reduce_large_data(): |
nothing calls this directly
no test coverage detected