mlx.data.Buffer.partition

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

Equivalent to slicing the buffer with a step equal to num_partitions and starting offset of partition.

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 buffer into.

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