PerLabelsPartitioner#

Per-labels dataset partitioning: shard-granularity interpolation from pathological skew to IID.

class krum.primitives.data_partitioners.per_labels.PerLabelsPartitioner[source]#

Bases: DataPartitioner

Per-labels partitioner: one sort-by-label mechanism spanning pathological skew to IID.

Sorts the dataset by label, cuts it into n_shards equal-size contiguous shards, shuffles the shard order, then deals shards to workers round-robin (worker \(w\) gets shards \(w, w+n, w+2n, \dots\)). Round-robin — rather than a fixed block of shards per worker — means any n_shards % n remainder shards are spread one at a time across the first few workers instead of being dropped, so worker dataset sizes never differ by more than one shard.

n_shards is itself controlled by \(\lambda \in [0, 1]\) (“iid-ness”), interpolating geometrically between the two extremes:

\[\text{n\_shards} = n \cdot \left(\frac{N}{n}\right)^{\lambda}\]

where \(N\) is the dataset size. At \(\lambda = 0\), n_shards = n — one giant, near-single-label shard per worker, the most pathological split this mechanism can produce. At \(\lambda = 1\), n_shards = N — every shard is a single sample, sorting becomes irrelevant at that granularity, and shuffle-then-round-robin reduces to exactly what IidPartitioner already does — recovering IID as a special case of the same mechanism, rather than needing a separate algorithm for it. Geometric (rather than linear) interpolation is deliberate: \(n\) and \(N\) typically span several orders of magnitude, and empirically, linear interpolation spends almost its entire range indistinguishable from IID, with the only interesting transition crammed into a tiny sliver near \(\lambda = 0\).

This is not a reproduction of any published scheme — it is an original design, built to unify IID and pathological sort-by-label skew under one continuously tunable mechanism instead of two separate partitioners.

classmethod partition(dataset: Dataset[Any], /, *, n: int, lambda_: float, seed: int = 42, **specialized: Any) → list[Subset[Any]][source]#

Split dataset across n workers via shard-granularity interpolation.

Parameters:

dataset – Full dataset to partition across workers. Labels are read from dataset.targets when available (as for the torchvision datasets), otherwise by indexing every sample.
n – Number of workers to split the dataset across.
lambda – Iid-ness in [0, 1]. 0 is the most pathological split this mechanism can produce (one shard per worker); 1 recovers plain IID (one sample per shard).
seed – Random seed for the shard-order shuffle.
**specialized – Additional keyword arguments (unused).

Returns:

List of ``n`` datasets, one per worker. Any remainder
(``len (dataset) % n_shards`` samples, past the last full shard)
is dropped, as in ``IidPartitioner``.

Raises:

ValueError – If n < 1, lambda_ is not in [0, 1], or dataset is nonempty but has fewer than n samples.