DataLoader

class braintools.trainer.DataLoader(dataset, batch_size=32, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, drop_last=False, prefetch=2, seed=None)

JAX-compatible data loader.

Parameters:

• dataset (Dataset | Any) – Dataset or array(s) to load from.

• batch_size (int) – Number of samples per batch.

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

• sampler (Sampler | None) – Custom sampler. Mutually exclusive with shuffle.

• batch_sampler (BatchSampler | None) – Custom batch sampler. Mutually exclusive with batch_size, shuffle, sampler.

• num_workers (int) – Number of worker processes for data loading (not implemented in JAX).

• collate_fn (Callable | None) – Function to collate samples into batches.

• drop_last (bool) – Whether to drop the last incomplete batch.

• prefetch (int) – Number of batches to prefetch (for device transfer).

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

Examples

Basic usage with arrays:

>>> X = jnp.ones((1000, 784))
>>> y = jnp.zeros((1000,))
>>> loader = DataLoader((X, y), batch_size=32, shuffle=True)
>>> for batch in loader:
...     x_batch, y_batch = batch
...     print(x_batch.shape)  # (32, 784)

With a dataset:

>>> dataset = DictDataset({'x': X, 'y': y})
>>> loader = DataLoader(dataset, batch_size=32)
>>> for batch in loader:
...     print(batch['x'].shape)  # (32, 784)

property num_samples: int

Return the total number of samples.

set_epoch(epoch)

Set the epoch (for distributed samplers).