---
id: grain
version: "0.2.18"
license: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# grain — Grain: A library for loading and transforming data for ML training.

License: permissive · Maintenance: active · Downloads: 2.1M/mo

## 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 above — 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

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

- License: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: medium
- Maintenance: active
- Downloads: 2.1M/month (top 5,000 on PyPI)
- Known vulnerabilities: none known

## Tags

data loading for machine learning, JAX data pipeline, batch processing library, deterministic data transformation, ML training data preparation, data-pipeline, machine-learning, jax

[View on SkillFed](https://skillfed.io/packages/grain) · [View on PyPI](https://pypi.org/project/grain/)