One function
EpochShuffledBatchSampler in gepa-ai/gepa
The author described this change as “Fix batch sampler so it updates when data-loader size increases and add tests to validate logic works”. It counts as a record because the checks below fail on the code as it stood at 7604b32be and pass on 2f825d627, with nothing else changed between the two runs.
Projectgepa-ai/gepa
Fix saved2025-10-23
Sharing licenceMIT · LICENSE
Change size+20 −3
What the code was meant to do, written into the code itself as a docstring
Mirrors the original batching logic: - Shuffle ids each epoch - Pad to minibatch size with least frequent ids - Deterministic via state.rng1
The change
| 11 | 11 | self.shuffled_ids: list[DataId] = [] | |
| 12 | 12 | self.epoch = -1 | |
| 13 | 13 | self.id_freqs = Counter() | |
| 14 | + | self.last_trainset_size = 0 | |
| 14 | 15 | if rng is None: | |
| 15 | 16 | self.rng = random.Random(0) | |
| 16 | 17 | else: | |
| ⋯ | |||
| 19 | 20 | def _update_shuffled(self, loader: DataLoader[DataId, DataInst]): | |
| 20 | 21 | all_ids = list(loader.all_ids()) | |
| 21 | 22 | trainset_size = len(loader) | |
| 23 | + | self.last_trainset_size = trainset_size | |
| 24 | + | ||
| 25 | + | if trainset_size == 0: | |
| 26 | + | self.shuffled_ids = [] | |
| 27 | + | self.id_freqs = Counter() | |
| 28 | + | return | |
| 29 | + | ||
| 22 | 30 | self.shuffled_ids = list(all_ids) | |
| 23 | 31 | self.rng.shuffle(self.shuffled_ids) | |
| 24 | - | for i in self.shuffled_ids: | |
| 25 | - | self.id_freqs[i] += 1 | |
| 32 | + | self.id_freqs = Counter(self.shuffled_ids) | |
| 26 | 33 | ||
| 27 | 34 | mod = trainset_size % self.minibatch_size | |
| 28 | 35 | num_to_pad = (self.minibatch_size - mod) if mod != 0 else 0 | |
| ⋯ | |||
| 33 | 40 | self.id_freqs[selected_id] += 1 | |
| 34 | 41 | ||
| 35 | 42 | def next_minibatch_ids(self, loader: DataLoader[DataId, DataInst], state: GEPAState) -> list[DataId]: | |
| 43 | + | trainset_size = len(loader) | |
| 44 | + | if trainset_size == 0: | |
| 45 | + | raise ValueError("Cannot sample a minibatch from an empty loader.") | |
| 46 | + | ||
| 36 | 47 | base_idx = state.i * self.minibatch_size | |
| 37 | 48 | curr_epoch = 0 if self.epoch == -1 else base_idx // max(len(self.shuffled_ids), 1) | |
| 38 | - | if curr_epoch > self.epoch: | |
| 49 | + | ||
| 50 | + | needs_refresh = ( | |
| 51 | + | not self.shuffled_ids | |
| 52 | + | or trainset_size != self.last_trainset_size | |
| 53 | + | or curr_epoch > self.epoch | |
| 54 | + | ) | |
| 55 | + | if needs_refresh: | |
| 39 | 56 | self.epoch = curr_epoch | |
| 40 | 57 | self._update_shuffled(loader) | |
| 41 | 58 | ||
The check that tells the two apart
fail→pass·tests/test_batch_sampler.py::test_epoch_sampler_errors_when_loader_empty
fail→pass·tests/test_batch_sampler.py::test_epoch_sampler_refreshes_when_loader_expands
Check file tests/test_batch_sampler.py, taken without changes from the fix and copied onto the older code, so the exact same check runs against both versions.
Origin and history
The code before it7604b32bea8e96b9ad594f6056075b484b680dba
Broken version dated2025-10-23
Modulegepa.strategies.batch_sampler
Units changedEpochShuffledBatchSampler
Fingerprint447fa5bae7a5125a
Checked2026-08-18 by goldset/0.1
Every field above is generated by our program. None of it is written by hand.
Other bugs found in gepa-ai/gepa
- 2026-08-11_from_legacy_config
- 2025-10-15Fix bug with merge overalp logic and add test case to cover correctness
- 2025-09-24fix: gepa state initialization from run_dir
- 2025-08-13Add tests and fix issue