MixingPartitioner#

Mixing dataset partitioning: interpolate between any two partitioners.

class krum.primitives.data_partitioners.mixing.MixingPartitioner[source]#

Bases: DataPartitioner

Mixing partitioner: interpolates between any two partitioners by a ratio.

Shuffles the dataset, splits it into a \((1 - \gamma)\) fraction and a \(\gamma\) fraction, partitions each fraction independently with p1 and p2 respectively (both across the same \(n\) workers), then gives worker \(w\) the concatenation of its p1-slice and its p2-slice.

\(\gamma = 0\) recovers p1 alone; \(\gamma = 1\) recovers p2 alone. This generalizes the “gamma-similarity” scheme of Karimireddy, Kale, Mohri, Reddi, Stich & Suresh (ICML 2020, SCAFFOLD, Section 7.1) — there, p1 is always an IID split and p2 is always a sort-by-label split — to any pair of partitioners.

Since the initial split accounts for every sample exactly once, and worker \(w\)’s dataset is the concatenation of its (disjoint) slice of each half, no sample is ever assigned to two workers. Whether a sample can be dropped depends on p1/p2 themselves: mixing in IidPartitioner (which drops a remainder within its own slice) can still drop samples, while DirichletPartitioner never does.

classmethod partition(dataset: Dataset[Any], /, *, n: int, p1: type[DataPartitioner], p2: type[DataPartitioner], gamma: float, p1_kwargs: dict[str, Any] | None = None, p2_kwargs: dict[str, Any] | None = None, seed: int = 42, **specialized: Any) → list[ConcatDataset[Any]][source]#

Split dataset across n workers by mixing p1 and p2.

Parameters:

dataset – Full dataset to partition across workers.
n – Number of workers to split the dataset across.
p1 – Partitioner applied to the \((1 - \gamma)\) fraction of the shuffled dataset.
p2 – Partitioner applied to the \(\gamma\) fraction.
gamma – Mixing ratio in [0, 1] — the fraction of the dataset routed to p2 instead of p1.
p1_kwargs – Extra keyword arguments forwarded to p1.partition (e.g. {"alpha": 0.5} when p1 is DirichletPartitioner).
p2_kwargs – Extra keyword arguments forwarded to p2.partition.
seed – Random seed for the initial shuffle, forwarded unchanged to both p1.partition and p2.partition.
**specialized – Additional keyword arguments (unused).

Returns:

List of ``n`` datasets, one per worker, each the concatenation
of that worker’s ``p1``-slice and ``p2``-slice.

Raises:

ValueError – If n < 1 or gamma is not in [0, 1].
TypeError – If p1 or p2 is not a DataPartitioner subclass.