grain

What it is and what it does

Grain is a data loading and transformation library designed for machine learning workflows. It provides a declarative API to define data processing pipelines—shuffling, mapping, batching, and other transformations—in a composable, deterministic way. While built with JAX models in mind, it does not require JAX and can work with other frameworks.

The library depends on absl-py, array-record, cloudpickle, etils, numpy, portpicker, and protobuf. It is actively maintained by Google, used in projects like MaxText and Gemma, and supports modern Python versions (3.11–3.14) across Linux, macOS, and Windows platforms.

Use it for:

• Define reproducible data pipelines for training JAX models with shuffling, mapping, and batching in a single declarative chain.
• Prepare and transform large datasets for machine learning experiments with deterministic, composable operations.
• Load and preprocess data for multi-framework ML workflows without being tied to a specific training framework.
• Build data augmentation and transformation steps that integrate seamlessly into ML training loops.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Grain is a Python library for reading, transforming, and batching data for training and evaluating machine learning models, with support for declarative data processing pipelines.

Yes. Grain is production-stable (Development Status 5), actively maintained, permissively licensed, and has no known vulnerabilities. Install friction is moderate but manageable. It is a good fit if you need declarative, deterministic data pipelines for ML training, especially with JAX, but also works with other frameworks.

Install

grain on PyPI

pip install grain

uv add grain

poetry add grain

Installing grain

Before you install

Medium install friction due to compiled wheels for multiple Python versions (3.11–3.14) and platforms. Active maintenance with recent release (58 days ago) and ongoing repository activity.

License in practice

Apache License 2.0 is permissive; you may use, modify, and distribute grain freely in commercial and private projects, provided you include license notices and document changes.

Quickstart

pip install grain

import grain

dataset = (
    grain.MapDataset.source([0, 1, 2, 3, 4, 5])
    .shuffle(seed=42)
    .map(lambda x: x + 1)
    .batch(batch_size=2)
)

for batch in dataset:
    print(batch)

Requires Python 3.11 or later; grain does not use GPU/TPU directly and runs transformations on CPU by default.

Verify before relying

• Whether grain's determinism guarantees hold across all transformation types and edge cases.
• Performance characteristics and scalability limits for very large datasets or complex pipelines.
• Compatibility with frameworks other than JAX beyond basic iteration.

Package facts

Evidence: grain-0.2.18-cp311-cp311-macosx_11_0_arm64.whl; grain-0.2.18-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; grain-0.2.18-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl; grain-0.2.18-cp311-cp311-win_amd64.whl; grain-0.2.18-cp312-cp312-macosx_11_0_arm64.whl; grain-0.2.18-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; grain-0.2.18-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl; grain-0.2.18-cp312-cp312-win_amd64.whl; grain-0.2.18-cp313-cp313-macosx_11_0_arm64.whl; grain-0.2.18-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; grain-0.2.18-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl; grain-0.2.18-cp313-cp313-win_amd64.whl; grain-0.2.18-cp314-cp314-macosx_11_0_arm64.whl; grain-0.2.18-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; grain-0.2.18-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl; grain-0.2.18-cp314-cp314-win_amd64.whl

Tags

More Artificial Intelligence packages