(inflated, baseline, batch_size, layout=None, oversized_shape=False)
| 46 | |
| 47 | |
| 48 | def check_batch(inflated, baseline, batch_size, layout=None, oversized_shape=False): |
| 49 | layout = layout or "" |
| 50 | assert inflated.layout() == layout, ( |
| 51 | f"The batch layout '({inflated.layout()})' does " |
| 52 | f"not match the expected layout ({layout})" |
| 53 | ) |
| 54 | inflated_samples = [np.array(sample) for sample in inflated.as_cpu()] |
| 55 | baseline_samples = [np.array(sample) for sample in baseline] |
| 56 | assert batch_size == len(inflated) == len(baseline) |
| 57 | if not oversized_shape: |
| 58 | for inflated_sample, baseline_sample in zip(inflated_samples, baseline_samples): |
| 59 | np.testing.assert_array_equal(inflated_sample, baseline_sample) |
| 60 | else: |
| 61 | for inflated_sample, baseline_sample in zip(inflated_samples, baseline_samples): |
| 62 | assert len(inflated_sample) == len(baseline_sample) |
| 63 | for inflated_frame, baseline_frame in zip(inflated_sample, baseline_sample): |
| 64 | flat_inflated = inflated_frame.reshape(-1) |
| 65 | baseline_size = baseline_frame.size |
| 66 | actually_inflated = flat_inflated[:baseline_size].reshape(baseline_frame.shape) |
| 67 | np.testing.assert_array_equal(actually_inflated, baseline_frame) |
| 68 | output_tail = flat_inflated[baseline_size:] |
| 69 | assert np.all(output_tail == 0), ( |
| 70 | f"Oversized output was not properly padded with 0s. " |
| 71 | f"Tail size {len(output_tail)}, the tail {output_tail}" |
| 72 | ) |
| 73 | |
| 74 | |
| 75 | def _test_sample_inflate(op, batch_size, np_dtype, seed): |
no test coverage detected