Grain¶

tinax.grain provides deterministic process-aware training and evaluation pipelines, random-access sources, and explicit multiprocessing worker lifetime.

Grain multiprocessing requires caller-parsed Abseil flags. Tinax validates this requirement but never mutates global flags.

InMemorySource ¶

InMemorySource(records: Iterable[T], *, source_id: str)

Snapshot an outer iterable while preserving references to its possibly mutable records.

A deterministic random-access source supporting len() and integer indexing. The outer iterable is copied into a tuple at construction, but individual records are held by reference and are not deep copied.

Parameters:

Name	Type	Description	Default
`records`	`Iterable[T]`	Iterable of records to snapshot into the source.	required
`source_id`	`str`	Stable, non-empty identifier used in source representations.	required

Raises:

Type	Description
`TypeError`	If `source_id` is not a string.
`ValueError`	If `source_id` is empty.

source_id `property` ¶

source_id: str

Return the caller-assigned stable identity used in source representations.

training_batches ¶

training_batches(
    dataset: MapDataset[T],
    *,
    shard_options: ShardOptions,
    seed: int,
    shuffle: bool,
    num_epochs: int | None,
    batch_size: int,
    batch_fn: BatchTransform[T, BatchT],
) -> MapDataset[BatchT]

Build deterministic full local batches with equal process step counts in every training epoch.

Parameters:

Name	Type	Description	Default
`dataset`	`MapDataset[T]`	Finite `grain.MapDataset` describing one epoch of records.	required
`shard_options`	`ShardOptions`	Per-process shard selection. `drop_remainder` must be `True` so every process yields the same number of steps.	required
`seed`	`int`	Integer seed for deterministic ordering and shuffling.	required
`shuffle`	`bool`	Whether to shuffle records within each epoch.	required
`num_epochs`	`int \| None`	Number of epochs to repeat, or `None` for an unbounded stream.	required
`batch_size`	`int`	Local (per-process) batch size. Must be positive.	required
`batch_fn`	`BatchTransform[T, BatchT]`	Callable collating a sequence of records into one batch.	required

Returns:

Type	Description
`MapDataset[BatchT]`	A `grain.MapDataset` of collated batches, sharded to this process and
`MapDataset[BatchT]`	repeated for `num_epochs` with per-epoch reseeding.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.MapDataset`, `shard_options` is not a `grain.sharding.ShardOptions`, `seed` is not an integer, `shuffle` is not a bool, `batch_fn` is not callable, or a size is a boolean.
`ValueError`	If `shard_options.drop_remainder` is not `True`, a size is not positive, the dataset is not one finite epoch, or it holds fewer than one complete global batch.

evaluation_batches ¶

evaluation_batches(
    dataset: MapDataset[T],
    *,
    shard_options: ShardOptions,
    batch_size: int,
    batch_fn: BatchTransform[T, BatchT],
) -> MapDataset[BatchT]

Build one ordered finite pass that keeps remainder records and may have unequal process step counts.

Parameters:

Name	Type	Description	Default
`dataset`	`MapDataset[T]`	Finite `grain.MapDataset` describing the evaluation records.	required
`shard_options`	`ShardOptions`	Per-process shard selection. `drop_remainder` must be `False` so remainder records are kept.	required
`batch_size`	`int`	Local (per-process) batch size. Must be positive.	required
`batch_fn`	`BatchTransform[T, BatchT]`	Callable collating a sequence of records into one batch.	required

Returns:

Type	Description
`MapDataset[BatchT]`	A `grain.MapDataset` of collated batches for this process's shard, in order
`MapDataset[BatchT]`	and without dropping the final partial batch.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.MapDataset`, `shard_options` is not a `grain.sharding.ShardOptions`, `batch_fn` is not callable, or `batch_size` is a boolean.
`ValueError`	If `shard_options.drop_remainder` is not `False`, `batch_size` is not positive, or the dataset is not one finite epoch.

open_multiprocessing_iterator ¶

open_multiprocessing_iterator(
    dataset: IterDataset[T],
    *,
    num_workers: int,
    per_worker_buffer_size: int = 1,
    worker_init_fn: Callable[[int, int], None]
    | None = None,
    sequential_slice: bool = False,
) -> Iterator[DatasetIterator[T]]

Open a Grain multiprocessing iterator and close every worker when the context exits.

A context manager yielding a dataset iterator backed by worker processes; on exit the iterator and all workers are closed.

Parameters:

Name	Type	Description	Default
`dataset`	`IterDataset[T]`	`grain.IterDataset` to iterate with prefetching workers.	required
`num_workers`	`int`	Number of worker processes. `0` runs in the main process.	required
`per_worker_buffer_size`	`int`	Prefetch buffer size per worker. Must be positive.	`1`
`worker_init_fn`	`Callable[[int, int], None] \| None`	Optional callable run in each worker as `worker_init_fn(worker_index, worker_count)`.	`None`
`sequential_slice`	`bool`	Whether workers slice the dataset sequentially.	`False`

Yields:

Type	Description
`DatasetIterator[T]`	A `grain.DatasetIterator` over the prefetched dataset.

Raises:

Type	Description
`TypeError`	If `dataset` is not a `grain.IterDataset`, an integer argument is a boolean, `worker_init_fn` is not callable or `None`, or `sequential_slice` is not a bool.
`ValueError`	If `num_workers` is negative or `per_worker_buffer_size` is not positive.
`RuntimeError`	If `num_workers` is nonzero but Abseil flags are not parsed; enter through `absl.app.run`.

Grain¶

InMemorySource ¶

source_id property ¶

training_batches ¶

evaluation_batches ¶

open_multiprocessing_iterator ¶

source_id `property` ¶