braintools.trainer.create_distributed_batches

braintools.trainer.create_distributed_batches(data, batch_size, devices=None, shuffle=False, seed=None)

Create batches suitable for distributed training with pmap.

Parameters:

• data (Any) – Input data (array, tuple of arrays, or dict of arrays).

• batch_size (int) – Batch size per device.

• devices (List[Any] | None) – List of devices. Default is all JAX devices.

• shuffle (bool) – Whether to shuffle the data.

• seed (int | None) – Random seed for shuffling.

Yields:

Any – Batches with shape (num_devices, batch_size, …).

Notes

Only full blocks of num_devices * batch_size samples are yielded; any trailing samples that cannot fill a complete block are dropped (this keeps the leading device axis a fixed size for pmap). Size the data or batch_size accordingly if you cannot afford to drop the remainder.

Examples

>>> for batch in create_distributed_batches(X, batch_size=32):
...     # batch.shape == (num_devices, 32, ...)
...     outputs = pmap_train_step(batch)