Transforms & Packing

Composable dataset transformations, batching, and sequence packing.

class taktiny.data.Compose(*functions)

Bases: Map

Apply deterministic record functions from left to right.

Accepts callables and Grain MapTransforms. With no functions this is an identity. Works directly on an example or as a DataLoader operation, and never iterates over examples implicitly. Use Filter, RandomMap, and Batch as separate loader operations, not inside Compose.

Parameters:

*functions (Callable[[Any], Any] | Map) – Ordered callables or Grain MapTransforms. Each receives the preceding function’s result. An empty composition is an identity.

Example

>>> from taktiny.data import Compose
>>> Compose(lambda x: x + 1, lambda x: x * 2)(3)
8

class taktiny.data.Map(function)

Bases: Map

Apply a callable to one record and return its replacement.

Records may be arrays, mappings, or arbitrary objects. The function decides what fields survive; Map does not merge its result with the original input. Use directly as a callable or as a DataLoader operation. Functions should avoid mutating source records and be serializable when using worker processes.

Parameters:

function (Callable[[Any], Any]) – Callable receiving one record and returning one output record.

Example

>>> from taktiny.data import DataLoader, Map
>>> operation = Map(lambda row: {'value': row['value'] * 2})
>>> operation({'value': 3})
{'value': 6}
>>> list(DataLoader([{'value': 1}, {'value': 2}], operations=[operation]))
[{'value': 2}, {'value': 4}]

class taktiny.data.MapFields(functions)

Bases: Map

Transform selected top-level mapping values, preserving other fields.

Each function receives its field’s value, not the complete record. Missing keys raise KeyError. A new dictionary is returned; unselected values are shared, not deep-copied. Functions should avoid mutating their inputs. Use Map for nested structures, renaming, or computations across fields.

Parameters:

functions (Mapping[Any, Callable[[Any], Any]]) – Mapping from top-level field keys to value-transforming callables. The mapping is copied at construction. An empty mapping returns a shallow copy of each input record.

Example

>>> from taktiny.data import MapFields
>>> MapFields({'value': lambda x: x / 255})({'value': 255, 'label': 2})
{'value': 1.0, 'label': 2}

class taktiny.data.RandomMap(function)

Bases: RandomMap

Apply random augmentation using Grain’s per-record NumPy Generator.

The loader’s sampler supplies the RNG. Use it instead of global random state for reproducible results and checkpoint restoration. With the default sampler, seed controls the stream and a new iterator starts it again. Direct calls require an explicit NumPy Generator, not a JAX random key.

Parameters:

function (Callable[[Any, Generator], Any]) – Callable receiving (record, rng) and returning one replacement record. Avoid mutating the original record.

Example

>>> from taktiny.data import DataLoader, RandomMap
>>> augment = RandomMap(lambda value, rng: value + int(rng.integers(10)))
>>> loader = DataLoader([1, 2, 3], operations=[augment], seed=42)
>>> first = list(loader)
>>> first == list(loader)
True

class taktiny.data.IndexMap(function)

Bases: MapWithIndex

Map records together with their sampler traversal indices.

The index is not necessarily the original source row key when shuffling, nor a dense output index after filtering. After batching, it belongs to the final contributing record. The default loader’s sharded sampler uses local traversal indices. Use map_with_index(index, element) for a direct call; __call__ adapts an iterator of Grain Records for DataLoader.

Parameters:

function (Callable[[int, Any], Any]) – Callable receiving (index, record) and returning a replacement record. Existing fields are preserved only if the function keeps them.

Example

>>> from taktiny.data import DataLoader, IndexMap
>>> operation = IndexMap(lambda index, value: {'index': index, 'value': value})
>>> list(DataLoader([10, 20], operations=[operation], shuffle=False))
[{'index': 0, 'value': 10}, {'index': 1, 'value': 20}]

class taktiny.data.Filter(function)

Bases: Filter

Keep records for which a predicate returns True.

Within DataLoader, rejected records are omitted without replacing or modifying accepted records. Calling Filter directly returns the predicate result, not the input record. Put it before Batch to filter individual rows.

Parameters:

function (Callable[[Any], bool]) – Callable receiving one record and returning a boolean.

Example

>>> from taktiny.data import DataLoader, Filter
>>> keep = Filter(lambda value: value >= 0)
>>> keep(-1)
False
>>> list(DataLoader([-1, 0, 2], operations=[keep]))
[0, 2]

class taktiny.data.FlatMap(function, *, max_fan_out)

Bases: FlatMap

Expand each example into zero or more examples, up to max_fan_out.

function(element) returns an iterable of records, not a column mapping. Expansion is bounded so Grain can checkpoint within an expanded record. Use this for windows, patches, segments, or any cardinality-changing map.

Parameters:

• function (Callable[[Any], Iterable[Any]]) – Callable receiving one record and returning an iterable of output records. Strings, bytes, and column mappings are not accepted as that iterable; wrap such records in a list instead.

• max_fan_out (int) – Positive upper bound on outputs per input record. Exceeding it raises ValueError. An empty output iterable drops the input.

Example

>>> from taktiny.data import DataLoader, FlatMap
>>> split = FlatMap(lambda values: values, max_fan_out=3)
>>> list(DataLoader([[1, 2], [], [3]], operations=[split]))
[1, 2, 3]
>>> split.flat_map([4, 5])
(4, 5)

class taktiny.data.Batch(batch_size, *, drop_remainder=False, collate_fn=None)

Bases: Batch

Group consecutive records, optionally using collate_fn(rows).

The default stacks matching leaves into NumPy arrays, preserving nested structure. It does not pad ragged data. Pass collate_fn=list to retain raw rows, or a custom callable for padding, audio, images, or arbitrary objects. With multiple workers, batching occurs within each worker independently.

Parameters:

• batch_size (int) – Positive maximum number of consecutive records in a batch.

• drop_remainder (bool) – If True, discard an incomplete final batch. Defaults to False, which emits it with a smaller leading batch dimension.

• collate_fn (Callable[[Sequence[Any]], Any] | None) – Optional callable receiving a sequence of rows and returning any batch structure. None uses Grain’s default leaf-wise stacking.

Example

>>> from taktiny.data import Batch, DataLoader
>>> loader = DataLoader([1, 2, 3], operations=[Batch(2)])
>>> [batch.tolist() for batch in loader]
[[1, 2], [3]]
>>> ragged = DataLoader([[1], [2, 3]], operations=[Batch(2, collate_fn=list)])
>>> list(ragged)
[[[1], [2, 3]]]

class taktiny.data.BatchMap(function, batch_size, *, drop_remainder=False)

Bases: object

Apply one callable to buffered rows and emit rows individually.

BatchMap expands into native Grain batch, map, and flat-map transformations. This preserves the cursor within a mapped batch when a dataloader iterator is checkpointed.

Unlike Batch, the output remains a stream of individual records. The function must preserve row count and order. Use FlatMap when changing cardinality; use Batch for final training batches.

Parameters:

• function (Callable[[Sequence[Any]], Any]) – Callable receiving a sequence of raw rows, not stacked columns. Return either one output per row, or a mapping of columns whose values each have the same length as the input buffer.

• batch_size (int) – Positive maximum number of input rows per function call.

• drop_remainder (bool) – Drop an incomplete final input buffer when True. With the default False, the function also receives the smaller buffer.

Example

>>> from taktiny.data import BatchMap, DataLoader
>>> operation = BatchMap(lambda rows: {'value': [x * 2 for x in rows]}, 2)
>>> list(DataLoader([1, 2, 3], operations=[operation]))
[{'value': 2}, {'value': 4}, {'value': 6}]

class taktiny.data.Pack(length, *, keys, axis=0, padding_values=None, position_key=None, mask_key=None, overflow='split', drop_remainder=False)

Bases: object

Concatenate aligned array fields into fixed-length records.

Parameters:

• length (int) – Positive number of steps along each field’s packing axis.

• keys (str | Sequence[str]) – One field name or a nonempty sequence of field names to pack. Only these fields and requested generated fields are emitted; other input fields are discarded because multiple source records can share one pack.

• axis (int | Mapping[str, int]) – Integer packing axis for all fields, or a mapping containing one axis per key. Negative axes count from the end of each array.

• padding_values (Mapping[str, Any] | None) – Per-field scalar padding values; all default to zero. Padding is cast to the field’s dtype, which is preserved.

• position_key (str | None) – Optional output name for a 1-D int32 vector of positions within each source record. Positions reset at record boundaries and continue across split fragments. Padding positions are zero.

• mask_key (str | None) – Optional name of a shared 1-D input/output validity mask. If present in an input, nonzero entries select steps from every field before packing. Missing input masks mean all steps are valid. Output masks are int32, one for data and zero for padding. Positions refer to the compressed sequence after mask selection.

• overflow (Literal['split', 'truncate']) – ‘split’ preserves all steps, carrying overflow into subsequent packs. ‘truncate’ fills the current pack and discards the rest of that source record, even if it could fill another pack.

• drop_remainder (bool) – Drop the final partial pack instead of right-padding it.

Fields must have equal lengths along their respective packing axes within each record. Each field must retain its dtype, rank, and non-packing shape across nonempty records. Different fields may have different dtypes/shapes. No token, label, or language-model conventions are applied. Inputs are not mutated; empty records are skipped. Arrays are processed with NumPy on CPU.

Use pack(records) for a lazy iterable of ordinary mapping records, or put this operation before Batch in a DataLoader. State belongs to each iterator, not the Pack object. In a loader, each worker/shard packs independently. Like the legacy PackSequences iterator operation, buffered packing state is not captured by Grain checkpoints: mid-pack resume is not supported.

Packing runs on individual records before the loader creates a batch. For a record shaped (time, channels), axis=0 packs time; final batching produces (batch, length, channels). No batch axis is assumed inside Pack.

Example

>>> import numpy as np
>>> from taktiny.data import DataLoader
>>> from taktiny.data.transforms import Pack
>>> packer = Pack(4, keys=('audio',), mask_key='valid')
>>> records = [{'audio': np.ones((3, 2), dtype=np.float32)}]
>>> result = next(packer.pack(records))
>>> result['audio'].shape
(4, 2)
>>> result['valid'].tolist()
[1, 1, 1, 0]
>>> loader = DataLoader(records, operations=[Pack(4, keys='audio')], batch_size=2)
>>> next(iter(loader))['audio'].shape
(1, 4, 2)

class taktiny.data.ApplyTemplate(template, *, format_fn=None, return_key='template')

Bases: Map

Format nested templates from mapping records and attach the result.

String leaves use str.format_map with the input record. Mapping values, lists, and tuples are formatted recursively; mapping keys and other leaves are kept as-is. A new record dictionary is returned, preserving other fields. Missing formatting fields raise KeyError. No tokenizer or text model is involved; templates can describe paths, metadata, or messages.

Parameters:

• template (Any) – String or nested structure to format. Containers are rebuilt for each record; arbitrary non-container leaves are shared.

• format_fn (Callable[[Any], Any] | None) – Optional keyword-only callable applied once to the complete formatted structure. None leaves it unchanged. The callable should avoid mutating shared template leaves.

• return_key (str) – Nonempty output field name; defaults to ‘template’. An existing field with this name is replaced in the returned record only.

Example

>>> from taktiny.data import ApplyTemplate
>>> operation = ApplyTemplate({'path': '{folder}/{name}.png'}, return_key='asset')
>>> operation({'folder': 'images', 'name': 'cat'})['asset']
{'path': 'images/cat.png'}
>>> join = ApplyTemplate(['{first}', '{last}'], format_fn=' '.join, return_key='name')
>>> join({'first': 'Ada', 'last': 'Lovelace'})['name']
'Ada Lovelace'