Skip to content

Sampling

Sampling primitives live in opake.dpsgd.sampling (Poisson, k-out-of-t allocation) and opake.dpftrl.sampling (cyclic Poisson with optional bands, b-min-sep, balls-in-bins, sequential). Distributed shard helpers live in opake.distributed. They provide privacy-amplifying sampling mechanisms for DP-SGD and DP-FTRL.

Overview

Privacy amplification through sampling is a key technique in DP-SGD: training on randomly selected subsets provides stronger privacy than training on the full dataset.

Opake provides these sampling strategies:

  1. Poisson Sampling — DP-SGD (opake.dpsgd.sampling.PoissonSampler): each example is sampled independently with probability sample_rate. Optional truncated_batch_size caps per-step batch size for more stable batches and memory; accounting must use the truncated-Poisson PLD (weaker than plain Poisson at the same sample_rate).

  2. K-Out-of-T Allocation — DP-SGD (opake.dpsgd.sampling.KOutOfTSampler): with allocation="block", each record is assigned to one batch in each of k contiguous, nearly equal blocks. With allocation="total", each record instead chooses k positions uniformly from the complete horizon.

  3. Cyclic Poisson (DP-FTRL) (opake.dpftrl.sampling.CyclicPoissonSampler): bands disjoint groups; step i samples only group i % bands, with independent inclusion at sample_rate. Use bands=1 for identity MF (full dataset each step); for BandMF, match bands to the MF strategy.

  4. Balls-in-Bins Sampling (BallsInBinsSampler): each example is independently assigned to a bin once at init; the assignment is fixed across epochs (required for BnB accounting). Bin sizes are variable (Binomial); some bins may be empty. Used with DP-λCGD, BISR, BSR, and BLT mechanisms.

  5. Sequential Batch Sampling (SequentialBatchSampler): iterates through the dataset in fixed-size contiguous batches with no randomness. The dataset should be pre-shuffled once before constructing the sampler. Used with the BLT mechanism.

  6. b-min-sep (BMinSepSampler): warm-start minimum-separation Poisson subsampling for BandMF (arXiv:2602.09338). Use with dpftrl_acc.b_min_sep.

See also: Sampling & Microbatching User Guide

PoissonSampler (DP-SGD)

from opake.dpsgd.sampling import PoissonSampler
from opake.random import key

sampler = PoissonSampler(
    data_source,
    sample_rate=batch_size / len(data_source),
    n_steps=None,
    truncated_batch_size=None,
    key=key(42),
)
loader = DataLoader(dataset, batch_sampler=sampler)
Parameter Type Default Description
data_source dataset with len() required The training dataset
sample_rate float required Probability of including each example, in (0, 1]. At sample_rate=1.0 there is no amplification — the accountant treats the step as the plain gaussian(nm)
n_steps int or None None Number of batches to yield. None = infinite
truncated_batch_size int or None None Optional per-step batch-size cap. When set, the sampler emits batches truncated to this many examples (uniform random subset of the Poisson draw).
key RngKey required RNG key for reproducible sampling

Plain Poisson — account with dpsgd_acc.poisson(dpsgd_acc.gaussian(nm), sample_rate).

Truncated Poisson (when truncated_batch_size is set) — account with dpsgd_acc.poisson(dpsgd_acc.gaussian(nm), sample_rate, truncated_batch_size=batch, dataset_size=n) to use the matching truncated-Poisson PLD. Truncated Poisson requires sample_rate < 1; for a full participation step (sample_rate=1.0) account directly with dpsgd_acc.gaussian(nm).

KOutOfTSampler (DP-SGD)

from opake.dpsgd.sampling import KOutOfTSampler
from opake.random import key

sampler = KOutOfTSampler(
    data_source,
    k=num_epochs,
    t=num_epochs * steps_per_epoch,
    allocation="block",
    key=key(42),
)
loader = DataLoader(dataset, batch_sampler=sampler)
Parameter Type Default Description
data_source dataset with len() required The training dataset
k int required Participations per record
t int required Total number of batches to yield
allocation "block" or "total" required Fixed-block or total-horizon allocation
key RngKey required RNG key for reproducible sampling

With allocation="block", the block sizes differ by at most one. Each block's batches partition the dataset exactly, with an independent assignment in each block. Bin sizes are Binomial, so some batches are empty; they are emitted rather than compacted away, because dropping them changes the accounted participation schedule.

Account block allocation with dpsgd_acc.k_out_of_t(mechanism, k=k, t=t, allocation="block"). This accounts the complete declared horizon once. Total allocation uses the same factory with allocation="total" and currently receives the conservative block bound.

Not the same scheme as BallsInBinsSampler

opake.dpftrl.sampling.BallsInBinsSampler draws the bin assignment once and reuses it for every epoch, because the matrix-mechanism dominating pair needs a known separation between an example's participations. KOutOfTSampler(..., allocation="block") draws each block independently, which is valid only because DP-SGD noise is uncorrelated across steps — and is strictly better there. Pair each sampler only with its own accountant.

BallsInBinsSampler

from opake.dpftrl.sampling import BallsInBinsSampler
from opake.random import key

sampler = BallsInBinsSampler(
    data_source,
    num_bins=dataset_size // batch_size,
    n_steps=8 * (dataset_size // batch_size),
    key=key(42),
)
loader = DataLoader(dataset, batch_sampler=sampler)
Parameter Type Default Description
data_source dataset with len() required The training dataset
num_bins int required Number of bins per epoch (≥ 2). Typically dataset_size / batch_size
n_steps int or None None Total number of batches to yield. Must be a positive multiple of num_bins (per-bin participation count is n_steps // num_bins). None = infinite
key RngKey required RNG key for reproducible sampling

Bin sizes are variable (Binomial distribution). Assignments are fixed across epochs (required for BnB dominating-pair accounting). Empty bins are emitted as empty batches so every accounted bin slot executes.

Account with dpftrl_acc.balls_in_bins(mechanism, num_bins=num_bins, n_steps=n_steps) where mechanism is dpftrl_acc.mf_gaussian(nm, strategy) for lambda_cgd_strategy, bisr_strategy, blt_strategy, bsr_strategy, or dpftrl_acc.mf_gaussian(..., identity_strategy()).

CyclicPoissonSampler (DP-FTRL)

Partitions the dataset into bands groups and, at step i, draws only from group i % bands, with each eligible example included independently at sample_rate (Binomial batch size within the group). Identity MF uses bands=1; BandMF uses bands equal to the mechanism’s band count — both pair with dpftrl_acc.poisson.

from opake.dpftrl.sampling import CyclicPoissonSampler
from opake.random import key

sampler = CyclicPoissonSampler(
    data_source,
    sample_rate=0.5,
    bands=4,
    n_steps=1000,
    key=key(42),
)
loader = DataLoader(dataset, batch_sampler=sampler)
Parameter Type Default Description
data_source dataset with len() required The training dataset
sample_rate float required Probability of including each eligible example, in (0, 1]
bands int 1 Groups in the cycle. 1 = identity-style (full dataset each step). >1 = cyclic BandMF-style (match strategy bands).
n_steps int 1 Total batches to yield
partition_type PartitionType EQUAL_SPLIT How to partition: EQUAL_SPLIT (only used when bands > 1) or INDEPENDENT
key RngKey required RNG key for reproducible sampling

In distributed training, shard the dataset with local_shard() and pass a per-rank key via fold_in(key, rank). Best used with mf_gaussian_noise for correlated noise (DP-FTRL); account with dpftrl_acc.poisson(mechanism, sample_rate, n_steps=...). There is no batch-size cap on this sampler; dpftrl_acc.poisson matches uncapped Poisson draws only.

Distributed Helpers

local_shard

Partition a dataset for DDP training. Returns a Subset containing the contiguous shard for the given rank.

from opake.distributed import local_shard
import torch.distributed as dist

shard = local_shard(
    dataset,
    rank=dist.get_rank(),
    world_size=dist.get_world_size(),
)
sampler = PoissonSampler(shard, sample_rate=0.01, key=fold_in(key(42), rank))
loader = DataLoader(shard, batch_sampler=sampler)
Parameter Type Default Description
dataset dataset with len() required The full training dataset
rank int 0 Current device rank
world_size int 1 Total number of devices

Returns: torch.utils.data.Subset containing the local shard.

local_shard(dataset, *, rank=0, world_size=1)

Return the shard of dataset that belongs to rank.

Each rank gets a contiguous, non-overlapping slice; the last rank receives any remainder examples.

Parameters:

Name Type Description Default
dataset Sized

Any object with __len__ (e.g. a PyTorch Dataset).

 required
rank int

Rank of the current worker (0-indexed). Defaults to 0.

 0
world_size int

Total number of workers. Defaults to 1.

 1

Returns:

Type Description
Subset

A torch.utils.data.Subset view into dataset.
Source code in packages/opake-engine/src/opake/api/engine/distributed/_shard.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def local_shard(dataset: Sized, *, rank: int = 0, world_size: int = 1) -> Subset:
    """Return the shard of ``dataset`` that belongs to ``rank``.

    Each rank gets a contiguous, non-overlapping slice; the last rank receives
    any remainder examples.

    Args:
        dataset: Any object with ``__len__`` (e.g. a PyTorch ``Dataset``).
        rank: Rank of the current worker (0-indexed). Defaults to 0.
        world_size: Total number of workers. Defaults to 1.

    Returns:
        A ``torch.utils.data.Subset`` view into ``dataset``.
    """
    start, end = _local_shard_bounds(len(dataset), rank=rank, world_size=world_size)
    return Subset(dataset, range(start, end))

API Documentation

PoissonSampler

Bases: Sampler

Poisson sampler for privacy amplification (DP-SGD).

Each example is independently included with probability sample_rate. When truncated_batch_size is set, batches are capped at that size (uniform without replacement from the Poisson sample).

For distributed training, shard the dataset externally and pass a per-rank key via fold_in(key, rank)::

from opake.distributed import local_shard

shard = local_shard(dataset, rank=rank, world_size=world_size)
sampler = PoissonSampler(
    shard, sample_rate=0.01, key=fold_in(key(42), rank)
)

Parameters:

Name Type Description Default
data_source Sized

Dataset to sample from (any object with __len__).

 required
sample_rate float

Probability of including each example ∈ (0, 1]. At 1.0 every example is in every batch — no amplification; account the step with :func:opake.dpsgd.accounting.gaussian.

 required
n_steps int | None

Number of batches to yield. None yields indefinitely.

 None
truncated_batch_size int | None

Optional cap on per-step batch size (truncated Poisson; use matching accounting—privacy is weaker than uncapped Poisson at the same sample_rate).

 None
key RngKey

RNG key for reproducibility.

 required

Example::

from opake.random import key
sampler = PoissonSampler(
    dataset, sample_rate=0.01, n_steps=1000, key=key(42),
)
loader = DataLoader(dataset, batch_sampler=sampler)
Note
  • Batch sizes are variable (Poisson property).
  • Expected batch size: len(data_source) * sample_rate (uncapped).
  • Use with DataLoader's batch_sampler parameter.
  • Resume via :func:opake.serialization.state_dict / :func:opake.serialization.from_state_dict: state captures the domain-separated stream key and consumed-step count.
Source code in packages/opake-dpsgd/src/opake/api/dpsgd/sampling/_poisson.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
class PoissonSampler(Sampler):
    """Poisson sampler for privacy amplification (DP-SGD).

    Each example is independently included with probability
    ``sample_rate``.  When ``truncated_batch_size`` is set, batches are
    capped at that size (uniform without replacement from the Poisson
    sample).

    For distributed training, shard the dataset externally and pass a
    per-rank key via ``fold_in(key, rank)``::

        from opake.distributed import local_shard

        shard = local_shard(dataset, rank=rank, world_size=world_size)
        sampler = PoissonSampler(
            shard, sample_rate=0.01, key=fold_in(key(42), rank)
        )

    Args:
        data_source: Dataset to sample from (any object with ``__len__``).
        sample_rate: Probability of including each example ``∈ (0, 1]``.
            At ``1.0`` every example is in every batch — no amplification;
            account the step with :func:`opake.dpsgd.accounting.gaussian`.
        n_steps: Number of batches to yield. ``None`` yields indefinitely.
        truncated_batch_size: Optional cap on per-step batch size (truncated
            Poisson; use matching accounting—privacy is weaker than uncapped
            Poisson at the same ``sample_rate``).
        key: RNG key for reproducibility.

    Example::

        from opake.random import key
        sampler = PoissonSampler(
            dataset, sample_rate=0.01, n_steps=1000, key=key(42),
        )
        loader = DataLoader(dataset, batch_sampler=sampler)

    Note:
        - Batch sizes are variable (Poisson property).
        - Expected batch size: ``len(data_source) * sample_rate`` (uncapped).
        - Use with ``DataLoader``'s ``batch_sampler`` parameter.
        - Resume via :func:`opake.serialization.state_dict` /
          :func:`opake.serialization.from_state_dict`: state captures the
          domain-separated stream key and consumed-step count.
    """

    def __init__(
        self,
        data_source: Sized,
        sample_rate: float,
        n_steps: int | None = None,
        truncated_batch_size: int | None = None,
        *,
        key: RngKey,
    ):
        self._initialize(
            data_source,
            sample_rate,
            n_steps,
            truncated_batch_size,
            stream_key=fold_in(key, POISSON_STREAM_FOLD),
        )

    @classmethod
    def _from_stream_key(
        cls,
        data_source: Sized,
        sample_rate: float,
        n_steps: int | None = None,
        truncated_batch_size: int | None = None,
        *,
        stream_key: RngKey,
    ) -> Self:
        """Construct from an already domain-separated stream key."""
        sampler = object.__new__(cls)
        sampler._initialize(
            data_source,
            sample_rate,
            n_steps,
            truncated_batch_size,
            stream_key=stream_key,
        )
        return sampler

    def _initialize(
        self,
        data_source: Sized,
        sample_rate: float,
        n_steps: int | None,
        truncated_batch_size: int | None,
        *,
        stream_key: RngKey,
    ) -> None:
        super().__init__()

        if not 0 < sample_rate <= 1:
            raise ConfigurationError(
                *(f"sample_rate must be in (0, 1], got {sample_rate}",)
            )
        if n_steps is not None and n_steps < 1:
            raise ConfigurationError(*(f"n_steps must be >= 1 or None, got {n_steps}",))
        if truncated_batch_size is not None:
            if type(truncated_batch_size) is not int:
                raise InputTypeError(
                    *(
                        "truncated_batch_size must be an int, got "
                        f"{type(truncated_batch_size).__name__}",
                    )
                )
            if truncated_batch_size < 1:
                raise ConfigurationError(
                    *(f"truncated_batch_size must be >= 1, got {truncated_batch_size}",)
                )

        self.data_source: Sized = data_source
        self.sample_rate = sample_rate
        self.n_steps = n_steps
        self.truncated_batch_size = truncated_batch_size

        self._num_samples = len(data_source)
        self._stream_key = stream_key
        self.generator = np.random.default_rng(stream_key.seed)
        self._consumed = 0

    def _sample_step(self) -> list[int]:
        return _plain_poisson_step_indices(
            self.generator,
            self._num_samples,
            self.sample_rate,
            self.truncated_batch_size,
        )

    def __iter__(self) -> Iterator[list[int]]:
        """Yield variable-size batches as lists of indices.

        Each iteration samples the entire dataset using Poisson subsampling,
        optionally truncated.  ``self._consumed`` increments after each
        yielded batch so state-dict snapshots reflect the live position.
        """
        if self.n_steps is None:
            while True:
                batch = self._sample_step()
                self._consumed += 1
                yield batch
        else:
            for _ in range(self._consumed, self.n_steps):
                batch = self._sample_step()
                self._consumed += 1
                yield batch

    def __len__(self) -> int:
        """Return the number of batches the iterator will yield.

        After a partial run, this is ``n_steps - consumed`` so
        ``len(DataLoader(batch_sampler=sampler))`` reflects the batches
        actually remaining (not the original total).  ``__iter__``
        starts at ``self._consumed``; ``__len__`` matches.

        Raises:
            TypeError: If n_steps is None (infinite iteration).
        """
        if self.n_steps is None:
            raise InputTypeError(*("len() of unsized object (n_steps=None)",))
        return self.n_steps - self._consumed

    @property
    def consumed(self) -> int:
        """Number of batches yielded so far (resume cursor)."""
        return self._consumed

    @property
    def expected_batch_size(self) -> float:
        """Expected batch size = num_samples * sample_rate (before truncation)."""
        return self._num_samples * self.sample_rate

    @property
    def batch_size_variance(self) -> float:
        """Variance of batch size for Poisson sampling (before truncation)."""
        return self._num_samples * self.sample_rate * (1 - self.sample_rate)

consumed property

Number of batches yielded so far (resume cursor).

expected_batch_size property

Expected batch size = num_samples * sample_rate (before truncation).

batch_size_variance property

Variance of batch size for Poisson sampling (before truncation).

__iter__()

Yield variable-size batches as lists of indices.

Each iteration samples the entire dataset using Poisson subsampling, optionally truncated. self._consumed increments after each yielded batch so state-dict snapshots reflect the live position.

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/sampling/_poisson.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def __iter__(self) -> Iterator[list[int]]:
    """Yield variable-size batches as lists of indices.

    Each iteration samples the entire dataset using Poisson subsampling,
    optionally truncated.  ``self._consumed`` increments after each
    yielded batch so state-dict snapshots reflect the live position.
    """
    if self.n_steps is None:
        while True:
            batch = self._sample_step()
            self._consumed += 1
            yield batch
    else:
        for _ in range(self._consumed, self.n_steps):
            batch = self._sample_step()
            self._consumed += 1
            yield batch

__len__()

Return the number of batches the iterator will yield.

After a partial run, this is n_steps - consumed so len(DataLoader(batch_sampler=sampler)) reflects the batches actually remaining (not the original total). __iter__ starts at self._consumed; __len__ matches.

Raises:

Type Description
TypeError

If n_steps is None (infinite iteration).
Source code in packages/opake-dpsgd/src/opake/api/dpsgd/sampling/_poisson.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def __len__(self) -> int:
    """Return the number of batches the iterator will yield.

    After a partial run, this is ``n_steps - consumed`` so
    ``len(DataLoader(batch_sampler=sampler))`` reflects the batches
    actually remaining (not the original total).  ``__iter__``
    starts at ``self._consumed``; ``__len__`` matches.

    Raises:
        TypeError: If n_steps is None (infinite iteration).
    """
    if self.n_steps is None:
        raise InputTypeError(*("len() of unsized object (n_steps=None)",))
    return self.n_steps - self._consumed

KOutOfTSampler

Bases: Sampler

Allocate every example to k of t training steps.

allocation="block" partitions the horizon into k contiguous, nearly equal blocks. Every example is assigned to one batch in each block, independently redrawing the assignment at block boundaries.

allocation="total" chooses exactly k distinct steps uniformly from the complete t-step horizon for every example. Its current accountant uses the block scheme as a conservative upper bound.

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/sampling/_k_out_of_t.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
class KOutOfTSampler(Sampler):
    """Allocate every example to ``k`` of ``t`` training steps.

    ``allocation="block"`` partitions the horizon into ``k`` contiguous,
    nearly equal blocks. Every example is assigned to one batch in each block,
    independently redrawing the assignment at block boundaries.

    ``allocation="total"`` chooses exactly ``k`` distinct steps uniformly from
    the complete ``t``-step horizon for every example. Its current accountant
    uses the block scheme as a conservative upper bound.
    """

    def __init__(
        self,
        data_source: Sized,
        *,
        k: int,
        t: int,
        allocation: _Allocation,
        key: RngKey,
    ):
        self._initialize(
            data_source,
            k=k,
            t=t,
            allocation=allocation,
            stream_key=fold_in(key, K_OUT_OF_T_STREAM_FOLD, allocation),
        )

    @classmethod
    def _from_stream_key(
        cls,
        data_source: Sized,
        *,
        k: int,
        t: int,
        allocation: _Allocation,
        stream_key: RngKey,
    ) -> Self:
        """Construct from an already domain-separated stream key."""
        sampler = object.__new__(cls)
        sampler._initialize(
            data_source,
            k=k,
            t=t,
            allocation=allocation,
            stream_key=stream_key,
        )
        return sampler

    def _initialize(
        self,
        data_source: Sized,
        *,
        k: int,
        t: int,
        allocation: _Allocation,
        stream_key: RngKey,
    ) -> None:
        super().__init__()
        if len(data_source) == 0:
            raise ConfigurationError(*("data_source must not be empty",))
        if t < 1:
            raise ConfigurationError(*(f"t must be >= 1, got {t}",))
        if not 1 <= k <= t:
            raise ConfigurationError(*(f"k must be in [1, t={t}], got {k}",))
        if allocation not in ("block", "total"):
            raise ConfigurationError(
                *(f"allocation must be 'block' or 'total', got {allocation!r}",)
            )

        self.data_source: Sized = data_source
        self.k = int(k)
        self.t = int(t)
        self.allocation = allocation
        self._num_samples = len(data_source)
        self._stream_key = stream_key
        self._consumed = 0

    @property
    def consumed(self) -> int:
        """Number of batches yielded so far (resume cursor)."""
        return self._consumed

    @property
    def expected_batch_size(self) -> float:
        """Horizon-average expected number of examples per batch."""
        return self._num_samples * self.k / self.t

    @property
    def block_sizes(self) -> tuple[int, ...] | None:
        """Contiguous block sizes, or ``None`` for total allocation."""
        if self.allocation == "total":
            return None
        floor = self.t // self.k
        num_ceil = self.t - floor * self.k
        return (floor,) * (self.k - num_ceil) + (floor + 1,) * num_ceil

    def _block_bins(self, block: int, size: int) -> list[list[int]]:
        """Return the independently drawn partition for one allocation block."""
        rng = np.random.default_rng(fold_in(self._stream_key, block).seed)
        assignment = rng.integers(0, size, size=self._num_samples)
        bins: list[list[int]] = [[] for _ in range(size)]
        for idx, bin_index in enumerate(assignment):
            bins[bin_index].append(idx)
        return bins

    def _total_batches(self) -> Iterator[list[int]]:
        """Yield the uniform global ``k``-out-of-``t`` allocation."""
        rng = np.random.default_rng(self._stream_key.seed)
        remaining = np.full(self._num_samples, self.k, dtype=np.int64)
        for step in range(self.t):
            probabilities = remaining / (self.t - step)
            mask = rng.random(self._num_samples) < probabilities
            yield np.flatnonzero(mask).tolist()
            remaining -= mask

    def __iter__(self) -> Iterator[list[int]]:
        if self.allocation == "block":
            floor = self.t // self.k
            num_ceil = self.t - floor * self.k
            num_floor = self.k - num_ceil
            floor_span = num_floor * floor
            bins: list[list[int]] | None = None
            current_block: int | None = None
            for step in range(self._consumed, self.t):
                if step < floor_span:
                    block, slot = divmod(step, floor)
                    size = floor
                else:
                    block_offset, slot = divmod(step - floor_span, floor + 1)
                    block = num_floor + block_offset
                    size = floor + 1
                if bins is None or block != current_block:
                    bins = self._block_bins(block, size)
                    current_block = block
                self._consumed = step + 1
                yield bins[slot]
            return

        for step, batch in enumerate(self._total_batches()):
            if step < self._consumed:
                continue
            self._consumed = step + 1
            yield batch

    def __len__(self) -> int:
        return self.t - self._consumed

consumed property

Number of batches yielded so far (resume cursor).

expected_batch_size property

Horizon-average expected number of examples per batch.

block_sizes property

Contiguous block sizes, or None for total allocation.

CyclicPoissonSampler

Bases: Sampler

Cyclic Poisson subsampling for DP-FTRL (one active group per step).

Disjoint example groups are fixed at construction (modulo partition_type). Step i samples only from group i % bands; each eligible example is a Bernoulli(sample_rate) draw, so batch size is random within that group and the active group advances each step.

Use bands=1 for identity MF (full dataset, plain Poisson each step, identity_strategy / identity_mf and ftrl_acc.poisson). For BandMF, set bands to the band count in band_mf_strategy / BandMf.

Parameters:

Name Type Description Default
data_source Sized

Dataset to sample from (any object with __len__).

 required
sample_rate float

Per-step inclusion probability ∈ (0, 1]. At 1.0 every example of the active group participates — no amplification; accounting treats each step as the plain Gaussian.

 required
bands int

Number of groups in the cycle. 1 = identity-style plain Poisson on the full dataset every step.

 1
n_steps int

Total number of batches to yield. Defaults to 1.

 1
partition_type PartitionType

Partition strategy when bands > 1.

 EQUAL_SPLIT
truncated_batch_size int | None

Optional cap on per-step batch size (truncated Poisson; use matching accounting—privacy is weaker than uncapped Poisson at the same sample_rate). Only the bands == 1 / IdentityMf combination is supported in :func:opake.dpftrl.accounting.poisson.

 None
key RngKey

RNG key for reproducibility.

 required

Example::

from opake.random import key
sampler = CyclicPoissonSampler(
    dataset, sample_rate=0.01, bands=4, n_steps=1000, key=key(42),
)
loader = DataLoader(dataset, batch_sampler=sampler)
Note

Batch sizes are variable (Poisson). Expected batch size per step is |group| * sample_rate where |group| = |D| / bands. Use with DataLoader's batch_sampler.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_poisson.py 
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
class CyclicPoissonSampler(Sampler):
    """Cyclic Poisson subsampling for DP-FTRL (one active group per step).

    Disjoint example groups are fixed at construction (modulo
    ``partition_type``).  Step ``i`` samples only from group ``i % bands``;
    each eligible example is a Bernoulli(``sample_rate``) draw, so batch size is
    random within that group and the active group advances each step.

    Use ``bands=1`` for identity MF (full dataset, plain Poisson each step,
    ``identity_strategy`` / ``identity_mf`` and ``ftrl_acc.poisson``).  For
    BandMF, set ``bands`` to the band count in ``band_mf_strategy`` / ``BandMf``.

    Args:
        data_source: Dataset to sample from (any object with ``__len__``).
        sample_rate: Per-step inclusion probability ``∈ (0, 1]``.  At
            ``1.0`` every example of the active group participates — no
            amplification; accounting treats each step as the plain Gaussian.
        bands: Number of groups in the cycle.  ``1`` = identity-style plain
            Poisson on the full dataset every step.
        n_steps: Total number of batches to yield.  Defaults to ``1``.
        partition_type: Partition strategy when ``bands > 1``.
        truncated_batch_size: Optional cap on per-step batch size (truncated
            Poisson; use matching accounting—privacy is weaker than uncapped
            Poisson at the same ``sample_rate``).  Only the ``bands == 1`` /
            ``IdentityMf`` combination is supported in
            :func:`opake.dpftrl.accounting.poisson`.
        key: RNG key for reproducibility.

    Example::

        from opake.random import key
        sampler = CyclicPoissonSampler(
            dataset, sample_rate=0.01, bands=4, n_steps=1000, key=key(42),
        )
        loader = DataLoader(dataset, batch_sampler=sampler)

    Note:
        Batch sizes are variable (Poisson).  Expected batch size per step is
        ``|group| * sample_rate`` where ``|group| = |D| / bands``.  Use with
        ``DataLoader``'s ``batch_sampler``.
    """

    def __init__(
        self,
        data_source: Sized,
        sample_rate: float,
        bands: int = 1,
        n_steps: int = 1,
        partition_type: PartitionType = PartitionType.EQUAL_SPLIT,
        truncated_batch_size: int | None = None,
        *,
        key: RngKey,
    ):
        self._initialize(
            data_source,
            sample_rate,
            bands,
            n_steps,
            partition_type,
            truncated_batch_size,
            stream_key=fold_in(key, CYCLIC_POISSON_STREAM_FOLD),
        )

    @classmethod
    def _from_stream_key(
        cls,
        data_source: Sized,
        sample_rate: float,
        bands: int = 1,
        n_steps: int = 1,
        partition_type: PartitionType = PartitionType.EQUAL_SPLIT,
        truncated_batch_size: int | None = None,
        *,
        stream_key: RngKey,
    ) -> Self:
        """Construct from an already domain-separated stream key."""
        sampler = object.__new__(cls)
        sampler._initialize(
            data_source,
            sample_rate,
            bands,
            n_steps,
            partition_type,
            truncated_batch_size,
            stream_key=stream_key,
        )
        return sampler

    def _initialize(
        self,
        data_source: Sized,
        sample_rate: float,
        bands: int,
        n_steps: int,
        partition_type: PartitionType,
        truncated_batch_size: int | None,
        *,
        stream_key: RngKey,
    ) -> None:
        super().__init__()

        if len(data_source) == 0:
            raise ConfigurationError(*("data_source must not be empty",))
        if not 0 < sample_rate <= 1:
            raise ConfigurationError(
                *(f"sample_rate must be in (0, 1], got {sample_rate}",)
            )
        if bands < 1:
            raise ConfigurationError(*(f"bands must be >= 1, got {bands}",))
        if n_steps < 1:
            raise ConfigurationError(*(f"n_steps must be >= 1, got {n_steps}",))
        if truncated_batch_size is not None and truncated_batch_size < 1:
            raise ConfigurationError(
                *(f"truncated_batch_size must be >= 1, got {truncated_batch_size}",)
            )

        self.num_examples = len(data_source)
        self._stream_key = stream_key
        self.generator = np.random.default_rng(stream_key.seed)

        if partition_type == PartitionType.INDEPENDENT:
            partition_fn = _independent_partition
        elif partition_type == PartitionType.EQUAL_SPLIT:
            partition_fn = _equal_split_partition
        else:
            raise ConfigurationError(
                *(f"Unsupported partition_type: {partition_type}",)
            )

        dtype = np.min_scalar_type(-self.num_examples)
        self.partition = partition_fn(self.num_examples, bands, self.generator, dtype)

        self.data_source: Sized = data_source
        self.sample_rate = sample_rate
        self.bands = bands
        self.n_steps = n_steps
        self.partition_type = partition_type
        self.truncated_batch_size = truncated_batch_size
        self._consumed = 0

    def _sample_step(self, step_idx: int) -> list[int]:
        """Sample one batch at the given step index.

        The active group rotates with ``step_idx % bands``; the per-step
        generator advance is the same regardless of whether the result
        is yielded or discarded, so state-dict replay can use this
        helper to fast-forward the generator without yielding.
        """
        group_idx = step_idx % self.bands
        group = self.partition[group_idx]

        sample_size = self.generator.binomial(n=len(group), p=self.sample_rate)
        if self.truncated_batch_size is not None:
            sample_size = min(sample_size, self.truncated_batch_size)

        if sample_size > 0:
            batch = self.generator.choice(
                group, size=sample_size, replace=False, shuffle=False
            )
        else:
            batch = np.array([], dtype=group.dtype)

        return batch.tolist()

    def __iter__(self) -> Iterator[list[int]]:
        """Yield Poisson batches.

        For each step, samples from group ``step % bands`` with inclusion
        probability ``sample_rate`` per example, optionally capped at
        ``truncated_batch_size``.
        """
        for step in range(self._consumed, self.n_steps):
            batch = self._sample_step(step)
            self._consumed += 1
            yield batch

    def __len__(self) -> int:
        """Batches remaining (``n_steps - consumed``).

        After a partial run, this reports what ``__iter__`` will yield
        — not the original total — so ``len(DataLoader(...))`` matches
        the resumed iteration count.
        """
        return self.n_steps - self._consumed

    @property
    def consumed(self) -> int:
        """Number of batches yielded so far (resume cursor)."""
        return self._consumed

    @property
    def expected_batch_size(self) -> float:
        """Expected batch size: ``|group| * sample_rate`` (before truncation)."""
        avg_group_size = self.num_examples / self.bands
        return avg_group_size * self.sample_rate

    @property
    def batch_size_variance(self) -> float:
        """Variance of batch size (Poisson property; before truncation)."""
        avg_group_size = self.num_examples / self.bands
        return avg_group_size * self.sample_rate * (1 - self.sample_rate)

consumed property

Number of batches yielded so far (resume cursor).

expected_batch_size property

Expected batch size: |group| * sample_rate (before truncation).

batch_size_variance property

Variance of batch size (Poisson property; before truncation).

__iter__()

Yield Poisson batches.

For each step, samples from group step % bands with inclusion probability sample_rate per example, optionally capped at truncated_batch_size.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_poisson.py 
211
212
213
214
215
216
217
218
219
220
221
def __iter__(self) -> Iterator[list[int]]:
    """Yield Poisson batches.

    For each step, samples from group ``step % bands`` with inclusion
    probability ``sample_rate`` per example, optionally capped at
    ``truncated_batch_size``.
    """
    for step in range(self._consumed, self.n_steps):
        batch = self._sample_step(step)
        self._consumed += 1
        yield batch

__len__()

Batches remaining (n_steps - consumed).

After a partial run, this reports what __iter__ will yield — not the original total — so len(DataLoader(...)) matches the resumed iteration count.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_poisson.py 
223
224
225
226
227
228
229
230
def __len__(self) -> int:
    """Batches remaining (``n_steps - consumed``).

    After a partial run, this reports what ``__iter__`` will yield
    — not the original total — so ``len(DataLoader(...))`` matches
    the resumed iteration count.
    """
    return self.n_steps - self._consumed

BMinSepSampler

Bases: Sampler

Warm-start b-min-sep sampler for BandMF privacy amplification.

Each iteration draws a batch by including each eligible example independently with probability p. Eligibility excludes any example that appeared in one of the previous bands - 1 batches (Algorithm 2).

Parameters:

Name Type Description Default
data_source Sized

Dataset with __len__.

 required
bands int

Minimum-separation parameter b (same as BandMF bandwidth).

 required
sampling_prob float

Paper's p in (0, 1] (per-iteration inclusion probability for each eligible example).

 required
n_steps int

Number of batches to yield.

 required
key RngKey

RNG key for reproducibility.

 required
Note

Use batch_sampler=... in DataLoader. Batch sizes are random (Poisson). For the same expected batch size as Poisson with per-example rate p_0 and b bands, set sampling_prob = p_0 / (1 - p_0 * (bands - 1)) (bands==1 → p_0).

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_b_min_sep.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
class BMinSepSampler(Sampler):
    """Warm-start b-min-sep sampler for BandMF privacy amplification.

    Each iteration draws a batch by including each *eligible* example
    independently with probability ``p``. Eligibility excludes any example
    that appeared in one of the previous ``bands - 1`` batches (Algorithm 2).

    Args:
        data_source: Dataset with ``__len__``.
        bands: Minimum-separation parameter ``b`` (same as BandMF bandwidth).
        sampling_prob: Paper's ``p`` in ``(0, 1]`` (per-iteration inclusion
            probability for each eligible example).
        n_steps: Number of batches to yield.
        key: RNG key for reproducibility.

    Note:
        Use ``batch_sampler=...`` in ``DataLoader``. Batch sizes are random
        (Poisson). For the same expected batch size as Poisson with
        per-example rate ``p_0`` and ``b`` bands, set
        ``sampling_prob = p_0 / (1 - p_0 * (bands - 1))`` (``bands==1`` → ``p_0``).
    """

    def __init__(
        self,
        data_source: Sized,
        bands: int,
        sampling_prob: float,
        n_steps: int,
        *,
        key: RngKey,
    ):
        self._initialize(
            data_source,
            bands,
            sampling_prob,
            n_steps,
            stream_key=fold_in(key, B_MIN_SEP_STREAM_FOLD),
        )

    @classmethod
    def _from_stream_key(
        cls,
        data_source: Sized,
        bands: int,
        sampling_prob: float,
        n_steps: int,
        *,
        stream_key: RngKey,
    ) -> Self:
        """Construct from an already domain-separated stream key."""
        sampler = object.__new__(cls)
        sampler._initialize(
            data_source,
            bands,
            sampling_prob,
            n_steps,
            stream_key=stream_key,
        )
        return sampler

    def _initialize(
        self,
        data_source: Sized,
        bands: int,
        sampling_prob: float,
        n_steps: int,
        *,
        stream_key: RngKey,
    ) -> None:
        super().__init__()
        if len(data_source) == 0:
            raise ConfigurationError(*("data_source must not be empty",))
        if bands < 1:
            raise ConfigurationError(*(f"bands must be >= 1, got {bands}",))
        if not 0 < sampling_prob <= 1:
            raise ConfigurationError(
                *(f"sampling_prob must be in (0, 1], got {sampling_prob}",)
            )
        if n_steps < 1:
            raise ConfigurationError(*(f"n_steps must be >= 1, got {n_steps}",))

        self.num_examples = len(data_source)
        self.data_source: Sized = data_source
        self.bands = bands
        self.sampling_prob = sampling_prob
        self.n_steps = n_steps
        self._stream_key = stream_key
        self.generator = np.random.default_rng(stream_key.seed)

        # Warm-start initial cooldown (Algorithm 2, lines 1–2).
        self._bars_left = np.zeros(self.num_examples, dtype=np.int32)
        if bands > 1:
            p = sampling_prob
            denom = 1.0 + (bands - 1) * p
            for j in range(self.num_examples):
                if self.generator.random() * denom < 1.0:
                    self._bars_left[j] = 0
                else:
                    # Uniform on {1, …, b-1}
                    self._bars_left[j] = int(self.generator.integers(1, bands))

        self._recent: deque[set[int]] = deque(maxlen=max(0, bands - 1))
        self._consumed = 0

    def _sample_step(self) -> list[int]:
        """Advance the Markov chain by one step and return the next batch.

        Mutates ``self.generator``, ``self._bars_left``, and
        ``self._recent``.  State-dict replay calls this helper to
        fast-forward without yielding.
        """
        exclude: set[int] = set()
        for batch in self._recent:
            exclude.update(batch)

        eligible = [
            idx for idx in range(self.num_examples) if self._bars_left[idx] == 0
        ]
        eligible = [i for i in eligible if i not in exclude]

        n_elig = len(eligible)
        if n_elig == 0:
            batch_indices: list[int] = []
        else:
            sample_size = self.generator.binomial(n=n_elig, p=self.sampling_prob)
            if sample_size > 0:
                pos = self.generator.choice(
                    n_elig, size=sample_size, replace=False, shuffle=False
                )
                batch_indices = [eligible[int(k)] for k in pos]
            else:
                batch_indices = []

        batch_set = set(batch_indices)

        # Advance cooldown: one step closer to eligibility.
        mask = self._bars_left > 0
        self._bars_left[mask] -= 1

        # Selected examples cannot appear in the next (bands - 1) iterations.
        for idx in batch_indices:
            self._bars_left[idx] = max(self._bars_left[idx], self.bands - 1)

        self._recent.append(batch_set)
        return batch_indices

    def __iter__(self) -> Iterator[list[int]]:
        for _ in range(self._consumed, self.n_steps):
            batch = self._sample_step()
            self._consumed += 1
            yield batch

    def __len__(self) -> int:
        """Batches remaining (``n_steps - consumed``).

        After a partial run, reflects what ``__iter__`` will yield —
        so ``len(DataLoader(...))`` matches the resumed iteration count.
        """
        return self.n_steps - self._consumed

    @property
    def consumed(self) -> int:
        """Number of batches yielded so far (resume cursor)."""
        return self._consumed

    @property
    def expected_batch_size(self) -> float:
        """Approximate expected batch size once the chain has mixed."""
        p0 = self.sampling_prob / (1.0 + self.sampling_prob * max(0, self.bands - 1))
        return self.num_examples * p0

    @property
    def batch_size_variance(self) -> float:
        """Approximate variance (Poisson on eligible pool ~ full size)."""
        p = self.sampling_prob
        n = float(self.num_examples)
        return n * p * (1.0 - p)

consumed property

Number of batches yielded so far (resume cursor).

expected_batch_size property

Approximate expected batch size once the chain has mixed.

batch_size_variance property

Approximate variance (Poisson on eligible pool ~ full size).

__len__()

Batches remaining (n_steps - consumed).

After a partial run, reflects what __iter__ will yield — so len(DataLoader(...)) matches the resumed iteration count.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_b_min_sep.py 
182
183
184
185
186
187
188
def __len__(self) -> int:
    """Batches remaining (``n_steps - consumed``).

    After a partial run, reflects what ``__iter__`` will yield —
    so ``len(DataLoader(...))`` matches the resumed iteration count.
    """
    return self.n_steps - self._consumed

BallsInBinsSampler

Bases: Sampler

Balls-in-Bins sampler: random independent bin assignment, fixed across epochs.

Each example independently chooses one of num_bins bins uniformly at random (true BnB, Definition 3.1 of Choquette-Choo et al. 2024). Bin sizes are random — some bins may be larger or smaller than the expected size len(dataset) / num_bins, and bins can even be empty. Empty bin slots yield empty batches so every epoch retains exactly num_bins optimizer steps.

The assignment is generated once and reused every epoch — this is required for the BnB privacy accounting (Lemma 3.2) to be valid.

For distributed training, shard the dataset externally and pass a per-rank key via fold_in(key, rank).

Parameters:

Name Type Description Default
data_source Sized

Dataset to sample from (any object with __len__).

 required
num_bins int

Number of bins per epoch (b ≥ 2). Typically dataset_size / desired_batch_size.

 required
n_steps int | None

Total number of batches to yield. Must be a positive multiple of num_bins (per-bin participation count is n_steps // num_bins). None yields indefinitely.

 None
key RngKey

RNG key for reproducibility.

 required
Example

from opake.random import key sampler = BallsInBinsSampler( ... dataset, num_bins=100, n_steps=1000, key=key(42) ... ) loader = DataLoader(dataset, batch_sampler=sampler)

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_balls_in_bins.py 
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class BallsInBinsSampler(Sampler):
    """Balls-in-Bins sampler: random independent bin assignment, fixed across epochs.

    Each example independently chooses one of ``num_bins`` bins uniformly
    at random (true BnB, Definition 3.1 of Choquette-Choo et al. 2024).
    Bin sizes are random — some bins may be larger or smaller than the
    expected size ``len(dataset) / num_bins``, and bins can even be empty.
    Empty bin slots yield empty batches so every epoch retains exactly
    ``num_bins`` optimizer steps.

    The assignment is generated once and reused every epoch — this is
    required for the BnB privacy accounting (Lemma 3.2) to be valid.

    For distributed training, shard the dataset externally and pass a
    per-rank key via ``fold_in(key, rank)``.

    Args:
        data_source: Dataset to sample from (any object with ``__len__``).
        num_bins: Number of bins per epoch (b ≥ 2).  Typically
            ``dataset_size / desired_batch_size``.
        n_steps: Total number of batches to yield.  Must be a positive
            multiple of ``num_bins`` (per-bin participation count is
            ``n_steps // num_bins``).  ``None`` yields indefinitely.
        key: RNG key for reproducibility.

    Example:
        >>> from opake.random import key
        >>> sampler = BallsInBinsSampler(
        ...     dataset, num_bins=100, n_steps=1000, key=key(42)
        ... )
        >>> loader = DataLoader(dataset, batch_sampler=sampler)
    """

    def __init__(
        self,
        data_source: Sized,
        num_bins: int,
        n_steps: int | None = None,
        *,
        key: RngKey,
    ):
        self._initialize(
            data_source,
            num_bins,
            n_steps,
            stream_key=fold_in(key, BALLS_IN_BINS_STREAM_FOLD),
        )

    @classmethod
    def _from_stream_key(
        cls,
        data_source: Sized,
        num_bins: int,
        n_steps: int | None = None,
        *,
        stream_key: RngKey,
    ) -> Self:
        """Construct from an already domain-separated stream key."""
        sampler = object.__new__(cls)
        sampler._initialize(
            data_source,
            num_bins,
            n_steps,
            stream_key=stream_key,
        )
        return sampler

    def _initialize(
        self,
        data_source: Sized,
        num_bins: int,
        n_steps: int | None,
        *,
        stream_key: RngKey,
    ) -> None:
        super().__init__()

        if len(data_source) == 0:
            raise ConfigurationError(*("data_source must not be empty",))
        if num_bins < _MIN_NUM_BINS:
            raise ConfigurationError(*(f"num_bins must be >= 2, got {num_bins}",))
        if n_steps is not None:
            if n_steps < 1:
                raise ConfigurationError(
                    *(f"n_steps must be >= 1 or None, got {n_steps}",)
                )
            if n_steps % num_bins != 0:
                raise ConfigurationError(
                    *(
                        f"n_steps ({n_steps}) must be a positive multiple of "
                        f"num_bins ({num_bins}); BnB analysis assumes integer epochs.",
                    )
                )

        self.data_source: Sized = data_source
        self.num_bins = num_bins
        self.n_steps = n_steps

        self._num_samples = len(data_source)
        self._stream_key = stream_key

        generator = np.random.default_rng(stream_key.seed)
        # True BnB: each example independently picks a bin.
        assignments = generator.integers(0, num_bins, size=self._num_samples)
        self._bins: list[list[int]] = [[] for _ in range(num_bins)]
        for idx, b in enumerate(assignments):
            self._bins[b].append(idx)
        self._consumed = 0

    @property
    def num_epochs(self) -> int | None:
        """Per-bin participation count: ``n_steps // num_bins``."""
        return None if self.n_steps is None else self.n_steps // self.num_bins

    @property
    def consumed(self) -> int:
        """Number of batches yielded so far (resume cursor)."""
        return self._consumed

    def __iter__(self) -> Iterator[list[int]]:
        """Yield every bin slot, including empty batches, for ``num_epochs``.

        The same bin assignment is reused every epoch — the BnB
        privacy accounting (Lemma 3.2) requires the assignment be fixed across
        the run. Empty slots must remain in the stream: dropping them would
        make the executed step schedule differ from the accountant's
        ``num_bins``-slot epoch.
        """
        bins = self._bins
        if self.n_steps is None:
            i = self._consumed
            while True:
                # Increment before yield so a snapshot taken mid-iter
                # reports the count of batches actually emitted so far.
                self._consumed = i + 1
                yield bins[i % self.num_bins]
                i += 1
        else:
            for i in range(self._consumed, self.n_steps):
                self._consumed = i + 1
                yield bins[i % self.num_bins]

    def __len__(self) -> int:
        """Declared bin slots remaining.

        After a partial run, reflects what ``__iter__`` will yield, including
        empty batches, so ``len(DataLoader(...))`` matches the accountant's
        remaining schedule.

        Raises:
            TypeError: If n_steps is None (infinite iteration).
        """
        if self.n_steps is None:
            raise InputTypeError(*("len() of unsized object (n_steps=None)",))
        return self.n_steps - self._consumed

    @property
    def expected_batch_size(self) -> float:
        """Expected batch size: len(dataset) / num_bins."""
        return self._num_samples / self.num_bins

    @property
    def sample_rate(self) -> float:
        """Effective sampling rate: 1 / num_bins."""
        return 1.0 / self.num_bins

num_epochs property

Per-bin participation count: n_steps // num_bins.

consumed property

Number of batches yielded so far (resume cursor).

expected_batch_size property

Expected batch size: len(dataset) / num_bins.

sample_rate property

Effective sampling rate: 1 / num_bins.

__iter__()

Yield every bin slot, including empty batches, for num_epochs.

The same bin assignment is reused every epoch — the BnB privacy accounting (Lemma 3.2) requires the assignment be fixed across the run. Empty slots must remain in the stream: dropping them would make the executed step schedule differ from the accountant's num_bins-slot epoch.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_balls_in_bins.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def __iter__(self) -> Iterator[list[int]]:
    """Yield every bin slot, including empty batches, for ``num_epochs``.

    The same bin assignment is reused every epoch — the BnB
    privacy accounting (Lemma 3.2) requires the assignment be fixed across
    the run. Empty slots must remain in the stream: dropping them would
    make the executed step schedule differ from the accountant's
    ``num_bins``-slot epoch.
    """
    bins = self._bins
    if self.n_steps is None:
        i = self._consumed
        while True:
            # Increment before yield so a snapshot taken mid-iter
            # reports the count of batches actually emitted so far.
            self._consumed = i + 1
            yield bins[i % self.num_bins]
            i += 1
    else:
        for i in range(self._consumed, self.n_steps):
            self._consumed = i + 1
            yield bins[i % self.num_bins]

__len__()

Declared bin slots remaining.

After a partial run, reflects what __iter__ will yield, including empty batches, so len(DataLoader(...)) matches the accountant's remaining schedule.

Raises:

Type Description
TypeError

If n_steps is None (infinite iteration).
Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_balls_in_bins.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
def __len__(self) -> int:
    """Declared bin slots remaining.

    After a partial run, reflects what ``__iter__`` will yield, including
    empty batches, so ``len(DataLoader(...))`` matches the accountant's
    remaining schedule.

    Raises:
        TypeError: If n_steps is None (infinite iteration).
    """
    if self.n_steps is None:
        raise InputTypeError(*("len() of unsized object (n_steps=None)",))
    return self.n_steps - self._consumed

SequentialBatchSampler

from opake.dpftrl.sampling import SequentialBatchSampler

sampler = SequentialBatchSampler(
    data_source,
    batch_size=256,
)
loader = DataLoader(dataset, batch_sampler=sampler)
Parameter Type Default Description
data_source dataset with len() required The training dataset (must not be empty)
batch_size int required Exact number of examples per batch (≥ 1)

Batch size is deterministic and fixed. The last incomplete batch is dropped (like drop_last=True). This sampler has no RNG key — it is fully deterministic. Pre-shuffle the dataset before constructing the sampler.

Used with the BLT mechanism, which requires deterministic batch order with fixed separation between participations.

SequentialBatchSampler

Bases: Sampler

Deterministic fixed-size sequential batching.

Yields contiguous, non-overlapping chunks of batch_size indices [0, ..., B-1], [B, ..., 2B-1], .... The last chunk is dropped when smaller than batch_size.

With n_steps set, the fixed batch order cycles: pass 2 repeats the identical batches in the identical order, so across the run every example participates exactly n_steps // num_batches times with exactly num_batches steps between participations — the min_sep / max_participations contract the BLT accounting assumes, enforced by one sampler over the whole stream.

This sampler has no RNG key — it is fully deterministic. Call dataset.shuffle(seed=...) once beforehand to randomise which examples land in which batch.

Parameters:

Name Type Description Default
data_source Sized

Dataset to sample from (any object with __len__).

 required
batch_size int

Exact number of examples per batch (must be ≥ 1).

 required
n_steps int | None

Total number of batches to yield across the run. Must be a positive multiple of the per-pass batch count (len(data_source) // batch_size), so participation counts stay uniform. None yields a single pass.

 None
Example

sampler = SequentialBatchSampler(dataset, batch_size=256) loader = DataLoader(dataset, batch_sampler=sampler)

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_sequential.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
class SequentialBatchSampler(Sampler):
    """Deterministic fixed-size sequential batching.

    Yields contiguous, non-overlapping chunks of ``batch_size`` indices
    ``[0, ..., B-1], [B, ..., 2B-1], ...``.  The last chunk is dropped
    when smaller than ``batch_size``.

    With ``n_steps`` set, the fixed batch order cycles: pass 2 repeats
    the identical batches in the identical order, so across the run
    every example participates exactly ``n_steps // num_batches`` times
    with exactly ``num_batches`` steps between participations — the
    ``min_sep`` / ``max_participations`` contract the BLT accounting
    assumes, enforced by one sampler over the whole stream.

    This sampler has no RNG key — it is fully deterministic.
    Call ``dataset.shuffle(seed=...)`` once beforehand to randomise
    which examples land in which batch.

    Args:
        data_source: Dataset to sample from (any object with ``__len__``).
        batch_size: Exact number of examples per batch (must be ≥ 1).
        n_steps: Total number of batches to yield across the run.  Must
            be a positive multiple of the per-pass batch count
            (``len(data_source) // batch_size``), so participation
            counts stay uniform.  ``None`` yields a single pass.

    Example:
        >>> sampler = SequentialBatchSampler(dataset, batch_size=256)
        >>> loader = DataLoader(dataset, batch_sampler=sampler)
    """

    def __init__(
        self,
        data_source: Sized,
        batch_size: int,
        n_steps: int | None = None,
    ):
        super().__init__()

        if len(data_source) == 0:
            raise ConfigurationError(*("data_source must not be empty",))
        if batch_size < 1:
            raise ConfigurationError(*(f"batch_size must be >= 1, got {batch_size}",))

        self.data_source: Sized = data_source
        self._num_samples = len(data_source)
        self._batch_size = batch_size
        self._num_batches = self._num_samples // self._batch_size

        if n_steps is not None:
            if n_steps < 1:
                raise ConfigurationError(
                    *(f"n_steps must be >= 1 or None, got {n_steps}",)
                )
            if self._num_batches == 0:
                raise ConfigurationError(
                    *(
                        f"batch_size ({batch_size}) exceeds dataset size "
                        f"({self._num_samples}); no complete batch to cycle.",
                    )
                )
            if n_steps % self._num_batches != 0:
                raise ConfigurationError(
                    *(
                        f"n_steps ({n_steps}) must be a positive multiple of the "
                        f"per-pass batch count ({self._num_batches}); a partial "
                        "final pass would make participation counts non-uniform.",
                    )
                )
        self.n_steps = n_steps
        self._consumed = 0

    def __iter__(self) -> Iterator[list[int]]:
        """Yield fixed-size batches of contiguous indices.

        The last incomplete batch (if any) is dropped.  With ``n_steps``
        set, the fixed order cycles until ``n_steps`` batches have been
        yielded.  Iteration resumes from ``self._consumed`` so a loaded
        sampler continues at its saved cursor.
        """
        total = self._num_batches if self.n_steps is None else self.n_steps
        for i in range(self._consumed, total):
            start = (i % self._num_batches) * self._batch_size
            # Increment before yield so a snapshot taken mid-iter
            # reports the count of batches actually emitted so far.
            self._consumed = i + 1
            yield list(range(start, start + self._batch_size))

    def __len__(self) -> int:
        """Declared batches remaining.

        After a partial run, reflects what ``__iter__`` will yield —
        so ``len(DataLoader(...))`` matches the resumed iteration count.
        """
        total = self._num_batches if self.n_steps is None else self.n_steps
        return total - self._consumed

    @property
    def consumed(self) -> int:
        """Number of batches yielded so far (resume cursor)."""
        return self._consumed

    @property
    def expected_batch_size(self) -> float:
        """Batch size (exact, not statistical)."""
        return float(self._batch_size)

consumed property

Number of batches yielded so far (resume cursor).

expected_batch_size property

Batch size (exact, not statistical).

__iter__()

Yield fixed-size batches of contiguous indices.

The last incomplete batch (if any) is dropped. With n_steps set, the fixed order cycles until n_steps batches have been yielded. Iteration resumes from self._consumed so a loaded sampler continues at its saved cursor.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_sequential.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def __iter__(self) -> Iterator[list[int]]:
    """Yield fixed-size batches of contiguous indices.

    The last incomplete batch (if any) is dropped.  With ``n_steps``
    set, the fixed order cycles until ``n_steps`` batches have been
    yielded.  Iteration resumes from ``self._consumed`` so a loaded
    sampler continues at its saved cursor.
    """
    total = self._num_batches if self.n_steps is None else self.n_steps
    for i in range(self._consumed, total):
        start = (i % self._num_batches) * self._batch_size
        # Increment before yield so a snapshot taken mid-iter
        # reports the count of batches actually emitted so far.
        self._consumed = i + 1
        yield list(range(start, start + self._batch_size))

__len__()

Declared batches remaining.

After a partial run, reflects what __iter__ will yield — so len(DataLoader(...)) matches the resumed iteration count.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/sampling/_sequential.py 
108
109
110
111
112
113
114
115
def __len__(self) -> int:
    """Declared batches remaining.

    After a partial run, reflects what ``__iter__`` will yield —
    so ``len(DataLoader(...))`` matches the resumed iteration count.
    """
    total = self._num_batches if self.n_steps is None else self.n_steps
    return total - self._consumed