xnn.common.distributed.shard_indices#

xnn.common.distributed.shard_indices(n)[source]#

The dataset indices this rank scores: every WORLD_SIZE-th one.

A strided shard, range(RANK, n, WORLD_SIZE), so that ranks hold disjoint subsets whose union is exactly the dataset. Unlike torch.utils.data.DistributedSampler nothing is padded or repeated, so metrics gathered across ranks equal the single-process ones; the price is that shards may differ in length by one, which evaluation does not mind (training does, which is why the trainer keeps the sampler). Usable directly as a DataLoader sampler.

Parameters:: n (int) – Dataset length.
Returns:: The indices for this rank (the full range(n) when not distributed).
Return type:: range