Skip to content

EpochShuffledBatchSampler

gepa.strategies.batch_sampler.EpochShuffledBatchSampler(minibatch_size: int, rng: random.Random | None = None)

Bases: BatchSampler[DataId, DataInst]

Mirrors the original batching logic: - Shuffle ids each epoch - Pad to minibatch size with least frequent ids - Deterministic via state.rng1

Source code in gepa/strategies/batch_sampler.py 
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(self, minibatch_size: int, rng: random.Random | None = None):
    self.minibatch_size = minibatch_size
    self.shuffled_ids: list[DataId] = []
    self.epoch = -1
    self.id_freqs = Counter()
    self.last_trainset_size = 0
    self._current_iteration: int | None = None
    self._calls_in_iteration = 0
    if rng is None:
        self.rng = random.Random(0)
    else:
        self.rng = rng

Attributes

minibatch_size = minibatch_size instance-attribute

shuffled_ids: list[DataId] = [] instance-attribute

epoch = -1 instance-attribute

id_freqs = Counter() instance-attribute

last_trainset_size = 0 instance-attribute

rng = random.Random(0) instance-attribute

Methods:

next_minibatch_ids(loader: DataLoader[DataId, DataInst], state: GEPAState) -> list[DataId]

Source code in gepa/strategies/batch_sampler.py 
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def next_minibatch_ids(self, loader: DataLoader[DataId, DataInst], state: GEPAState) -> list[DataId]:
    trainset_size = len(loader)
    if trainset_size == 0:
        raise ValueError("Cannot sample a minibatch from an empty loader.")

    # Repeated calls within one iteration (multi-proposal sampling
    # strategies request one minibatch per task) advance one chunk each,
    # so tasks in the same iteration get distinct minibatches. The first
    # call of an iteration is unchanged from the classic behavior. Chunks
    # wrap around, so distinctness holds only while the iteration's call
    # count stays within len(shuffled_ids) / minibatch_size.
    if state.i == self._current_iteration:
        self._calls_in_iteration += 1
    else:
        self._current_iteration = state.i
        self._calls_in_iteration = 0

    base_idx = state.i * self.minibatch_size
    curr_epoch = 0 if self.epoch == -1 else base_idx // max(len(self.shuffled_ids), 1)

    needs_refresh = not self.shuffled_ids or trainset_size != self.last_trainset_size or curr_epoch > self.epoch
    if needs_refresh:
        self.epoch = curr_epoch
        self._update_shuffled(loader)

    assert len(self.shuffled_ids) >= self.minibatch_size
    assert len(self.shuffled_ids) % self.minibatch_size == 0

    # The epoch bookkeeping above uses the un-offset base_idx (constant
    # within an iteration), so repeat calls never trigger a reshuffle and
    # the shuffle sequence stays identical to the single-call path.
    base_idx = (base_idx + self._calls_in_iteration * self.minibatch_size) % len(self.shuffled_ids)
    end_idx = base_idx + self.minibatch_size
    assert end_idx <= len(self.shuffled_ids)
    return self.shuffled_ids[base_idx:end_idx]