| 1118 | |
| 1119 | |
| 1120 | class PerIterShapeSource: |
| 1121 | def __init__(self, shapes): |
| 1122 | self.shapes = shapes |
| 1123 | |
| 1124 | def __call__(self, sample_info): |
| 1125 | batch_idx = sample_info.iteration |
| 1126 | shape = self.shapes[batch_idx % len(self.shapes)] |
| 1127 | return np.full(shape, sample_info.idx_in_epoch, dtype=np.uint8) |
| 1128 | |
| 1129 | |
| 1130 | def per_iter_shape_pipeline( |