grain
Grain: A library for loading and transforming data for ML training.
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python 3.11 or later; grain does not use GPU/TPU directly and runs transformations on CPU by default.
- 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 · maintenance · safety
permissive license (permissive) — 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.
last release 2026-06-17 (58 days) · last repo commit 2026-08-12 · 766 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 2,127,414 downloads/mo, #3,270 on PyPI
Alternatives
Verify before relying
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)- 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.
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 on it.
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
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.
Requires Python 3.11 or later; grain does not use GPU/TPU directly and runs transformations on CPU by default.
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)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 | permissive license permissive |
| Python support | Supports the current Python release >=3.11 |
| Install friction | Medium. Platform-specific wheel |
| Runtime dependencies | 7 packagesabsl-pyarray-recordcloudpickleetilsnumpyportpickerprotobuf |
| Maintenance | Actively maintained 58 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 2,127,414 / month, #3,270 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 5 - Production/StableProgramming Language :: Python :: 3.11Programming Language :: Python :: 3.12Programming Language :: Python :: 3.13Programming Language :: Python :: 3.14 |
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
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “data loading for machine learning”
- grainGrain is a Python library for reading, transforming, and batching…
- tensorboard-data-serverProvides fast data loading and serving for TensorBoard, the web…
- datazetsDatazets provides a simple interface to download and import…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also jax · jaxlib · flax · litdata · feast · webdataset · datasets · spark-nlp · mmengine · clu