Data Loading

Data loader and dataset splitting utilities.

class taktiny.data.DataLoader(source, *, operations=(), batch_size=None, drop_remainder=False, collate_fn=None, sampler=None, shuffle=False, seed=0, num_epochs=1, shard_index=0, shard_count=1, worker_count=0, worker_buffer_size=1)[source]

Bases: DataLoader

Preprocess and iterate over caller-provided, random-access records.

Records may be arrays, mappings, tuples, images, audio, strings, or custom objects. The source is never downloaded, decoded, or copied into memory. Supply a list, array, already-loaded dataset, or an object implementing __len__ and integer __getitem__. This loader does not accept streaming generators; use a streaming backend directly, or explicitly materialize a finite stream with list(source) if it fits in memory.

Parameters:

  • source (RandomAccessDataSource | Any) – Caller-owned random-access data, not a repository ID or path.

  • operations (Sequence[Any]) – Ordered Taktiny or native Grain operations. Use Map for a record callable, RandomMap for augmentation, and Filter to drop records. Operations are lazy and run when the loader is iterated.

  • batch_size (int | None) – Optional final batch size, applied after all operations. None emits records unchanged. For operations after batching, put Batch directly in operations and leave this argument unset.

  • drop_remainder (bool) – Drop an incomplete final batch; requires batch_size.

  • collate_fn (Callable[[Sequence[Any]], Any] | None) – Optional function(rows) for final batching. None uses Grain stacking; list preserves ragged or custom objects without padding. Requires batch_size. Collation runs independently in each worker.

  • sampler (Sampler | None) – Optional native Grain sampler. When provided, it owns sampling, epochs, seed, and sharding; the corresponding convenience arguments are ignored after validation.

  • shuffle (bool) – Shuffle indices using seed; defaults to False.

  • seed (int) – Unsigned 32-bit integer seed for sampling and RandomMap augmentation.

  • num_epochs (int | None) – Positive epoch count (default 1); None repeats indefinitely. Each new iterator starts from the beginning unless state is restored.

  • shard_index (int) – This process’s data shard, in [0, shard_count).

  • shard_count (int) – Number of data shards. Shards may have unequal lengths; this is input partitioning, not JAX device-array sharding.

  • worker_count (int | None) – Child workers; 0 runs locally, None lets Grain choose. Sources and transforms must be serializable when workers are used.

  • worker_buffer_size (int) – Positive per-worker prefetch buffer size.

Iterators retain Grain’s get_state()/set_state() checkpoint API. Restore against the same source and pipeline. Custom iterator operations retain their own Grain checkpoint limitations. Transforms should be deterministic apart from RandomMap’s supplied RNG and should not mutate source records.

Example

>>> from taktiny.data import DataLoader, MapFields
>>> rows = [{'value': 255, 'label': 0}, {'value': 0, 'label': 1}]
>>> loader = DataLoader(rows, operations=[
...     MapFields({'value': lambda x: x / 255})], batch_size=2)
>>> next(iter(loader))['value'].tolist()
[1.0, 0.0]
__init__(source, *, operations=(), batch_size=None, drop_remainder=False, collate_fn=None, sampler=None, shuffle=False, seed=0, num_epochs=1, shard_index=0, shard_count=1, worker_count=0, worker_buffer_size=1)[source]

Create a Grain loader from a random-access dataset.

operations are applied exactly in the supplied order. Mapping, filtering, packing, batching, and collation therefore remain separate concerns and can be composed using Grain transformations or custom Grain operations.

When sampler is omitted, an grain.IndexSampler is created from the remaining sampling arguments. Supplying sampler transfers sampling and sharding responsibility entirely to that object. The default num_epochs=1 creates a single-epoch (finite) loader; pass None for an unbounded loader.

Parameters:

  • source (RandomAccessDataSource | Any)

  • operations (Sequence[Any])

  • batch_size (int | None)

  • drop_remainder (bool)

  • collate_fn (Callable[[Sequence[Any]], Any] | None)

  • sampler (Sampler | None)

  • shuffle (bool)

  • seed (int)

  • num_epochs (int | None)

  • shard_index (int)

  • shard_count (int)

  • worker_count (int | None)

  • worker_buffer_size (int)

Return type:

None

class taktiny.data.RandomAccessSource(*args, **kwargs)[source]

Bases: Protocol

Structural interface for finite sources; no framework inheritance needed.

taktiny.data.train_validation_split(source, validation_size, *, shuffle=True, seed=0)[source]

Split a random-access source into (train, validation) views.

validation_size may be a count (int) or a fraction (float in (0, 1)). The returned views are random-access and can be passed directly to DataLoader. Fractions are rounded to the nearest integer; both splits must be nonempty. The source itself is not copied or shuffled.

Return type:

tuple[Any, Any]

Parameters: