mlx.data.Stream.partition

Stream.partition(self: mlx.data._c.Stream, num_partitions: int, partition: int) → mlx.data._c.Stream

For every num_partitions consecutive samples return the partition-th.

This can be used for distributed settings where different nodes should load different parts of a dataset.

Parameters:

• num_partitions (int) – How many different partitions to split the stream into.

• partition (int) – Which partition to use (0-based).