Grain Documentation

Grain is an open-source library designed for reading data to train and evaluate JAX models. It is optimized for performance, determinism, and resilience in machine learning pipelines.

Key Features:

• Powerful: Supports arbitrary Python transformations.
• Flexible: Modular design allows users to override Grain components with custom implementations.
• Deterministic: Ensures multiple runs of the same pipeline produce identical output.
• Resilient to Preemptions: Designed with minimal checkpoint sizes, allowing Grain to resume from a preemption point and produce the same output as if no interruption occurred.
• Performant: Optimized for various data modalities including Text, Audio, Images, and Video.
• Minimal Dependencies: Avoids heavy dependencies (e.g., it does not depend on TensorFlow).

The grain.transforms module provides a suite of transformation classes used to manipulate datasets during training. These transformations can be composed to create complex data pipelines. Key transformation types include:

• Map: Applies a function to every element in the dataset.
• MapWithIndex: Applies a function to every element, providing the element's index.
• Filter: Removes elements from the dataset based on a predicate function.
• RandomMap: Applies a function to elements with a stochastic component (useful for data augmentation).
• Batch: Groups consecutive elements into batches of a specified size.
• Transformations: A container class used to manage and compose multiple transformations.
• DatasetSelectionMap: Maps dataset indices to a selection of the original dataset.
• Transformation: The base class for all transformation logic.

The grain.sources module provides various data source implementations designed to feed JAX models. These sources are typically used to provide structured data to a training pipeline. The available data source classes include:

• RandomAccessDataSource: A base or specific implementation for random access patterns.
• ArrayRecordDataSource: A data source optimized for reading ArrayRecord files.
• SharedMemoryDataSource: A data source utilizing shared memory for efficient data access.
• RangeDataSource: A data source that provides data over a specific range.

All these classes implement the standard Python sequence protocol, supporting __len__ (to get the total number of elements) and __getitem__ (to access specific elements by index).

The grain.samplers module provides classes for defining how data is sampled during training or evaluation. It includes implementations for sequential sampling, index-based sampling, and a base Sampler interface.

Available Sampler types:

• Sampler: The base class for all samplers.
• SequentialSampler: Samples elements in a fixed, sequential order.
• IndexSampler: Samples elements based on a provided set of indices.

All samplers implement the standard Python sequence protocol, supporting __len__ and __getitem__.

The grain.checkpoint module provides tools for saving and restoring model states (checkpoints) during training. It is composed of three primary classes:

1. CheckpointHandler: The main interface for managing checkpointing operations.
2. CheckpointSave: A component or strategy used to define how and what to save.
3. CheckpointRestore: A component or strategy used to define how to reload a saved state.

To implement checkpointing in your JAX training loop, you typically interact with the CheckpointHandler to orchestrate the saving and loading of state dictionaries.

Grain is designed for feeding JAX models by providing data loading pipelines. The package is organized into two main API categories:

1. Dataset APIs: Low-level building blocks for defining data structures and transformations.
2. DataLoader APIs: High-level abstractions for managing data iteration, often used for training steps.

Key subpackages include:

• grain.checkpoint: For managing state and checkpoints.
• grain.samplers: For defining how data is sampled.
• grain.sharding: For distributed data loading.
• grain.transforms: For data augmentation and processing.
• grain.sources: For defining data origins.

The grain.experimental module provides advanced data loading primitives and transformations for JAX models. It includes specialized dataset types for handling different file formats (Parquet, TFRecord), advanced shuffling strategies (WindowShuffle), and performance optimizations like prefetching and caching.

Key components include:

• Iterative Datasets: Classes like ParquetIterDataset, TFRecordIterDataset, and InterleaveIterDataset for streaming data.
• Transformations: Functions like apply_transformations and batch_and_pad to manipulate data streams.
• Performance Tools: ThreadPrefetchIterDataset for overlapping data loading with computation and CacheIterDataset for speeding up repeated passes.

JAX recompiles whenever input shapes change. To prevent this, pair .batch(drop_remainder=True) with .repeat() to create an infinite stream that never produces a short final batch.

ds = (
    grain.MapDataset.source(source)
    .seed(42)
    .shuffle()
    .repeat()  # infinite stream
    .map(lambda r: {"image": r["image"].astype(np.float32) / 255.0,
                    "label": r["label"]})
    .batch(128, drop_remainder=True)
)

The Dataset API is a lower-level interface that uses chaining syntax to define transformations. It is designed for complex pipelines and provides more control over execution. The API consists of three primary classes:

• MapDataset: A dataset that supports efficient random access. It behaves like an (infinite) Sequence that computes values lazily. Most pipelines begin with one or more MapDataset objects (often derived from a RandomAccessDataSource).
• IterDataset: A dataset that does not support efficient random access and only supports iteration (behaves like an Iterable). You can convert any MapDataset into an IterDataset by calling .to_iter_dataset().
• DatasetIterator: A stateful iterator for an IterDataset. The state of this iterator can be saved and restored, which is useful for resuming training.

The DataLoader is the central orchestrator that glues three main components together to create a data pipeline:

1. Sampler: Determines which records to read next. It handles global transformations like shuffling, repeating for multiple epochs, and sharding across machines. It produces metadata objects containing an index (for checkpointing), a record_key, and an rng (for random transformations).
2. DataSource: Responsible for the low-level reading of individual records from storage (e.g., ArrayRecordDataSource or tfds.data_source).
3. Transformations: Applied to the records read from the data source to produce the final output elements (e.g., MapTransform, BatchTransform).

The DataLoader manages worker processes to parallelize these operations and provides an iterator to consume the processed elements.