Install & Compatibility
Where this runs
tested against v0.2.13 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.95 runs
build_error
glibcpy 3.10–3.95 runs
installs and imports cleanly · install 5.5s · import 0.978s · 105MB
111MB installed
● package 111MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
MapDataset
✓ import grain
dataset = grain.MapDataset.source([...])
IterDataset
✓ import grain
iter_dataset = dataset.to_iter_dataset()
✗ from grain.python.experimental import MultiprocessPrefetchIterDataset
MultiprocessPrefetchIterDataset and ConcatenateMapDataset were deprecated in 0.2.16; use `IterDataset.mp_prefetch` or `MapDataset.concatenate` instead.
This example demonstrates how to create a simple `MapDataset` from a list, apply common transformations like shuffling, mapping, and batching, and then iterate through the processed data. It showcases the declarative chaining API for data pipeline construction.
import grain
dataset = (
grain.MapDataset.source([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
.shuffle(seed=42) # Shuffles elements globally.
.map(lambda x: x + 1) # Maps each element.
.batch(batch_size=2) # Batches consecutive elements.
)
print("Processing dataset:")
for batch in dataset:
print(batch)
Debug
Known issues
breakingCustom implementations of `RandomAccessDataSource` must now accept an `int` index in `__getitem__`. While legacy paths handling `SupportsIndex` still work at runtime, type checkers may flag errors. Switch to `int` for full compatibility.fixUpdate `__getitem__` methods in custom `RandomAccessDataSource` implementations to accept an `int` argument: `def __getitem__(self, index: int):`
affects: 0.2.16 and later
deprecatedSupport for Python 3.10 has been deprecated, and the library now requires Python >=3.11.fixUpgrade your Python environment to 3.11 or newer.
affects: 0.2.14 and later
deprecatedExperimental APIs `grain.python.experimental.MultiprocessPrefetchIterDataset` and `grain.python.experimental.ConcatenateMapDataset` have been deprecated. Use their graduated versions `grain.IterDataset.mp_prefetch` and `grain.MapDataset.concatenate` respectively.fixMigrate usage from `grain.python.experimental.MultiprocessPrefetchIterDataset` to `grain.IterDataset.mp_prefetch` and from `grain.python.experimental.ConcatenateMapDataset` to `grain.MapDataset.concatenate`.
affects: 0.2.16 and later
gotchaWhen using Python multiprocessing for parallel data loading and transformations, all custom transformation functions (e.g., `MapTransform` subclasses) must be picklable. Non-picklable objects or closures can lead to errors during serialization.fixEnsure that any custom transformations are defined at the top level of a module or as static methods/free functions, and avoid using complex closures or unpicklable objects within them.
affects: All versions
gotchaChoose between `MapDataset` and `IterDataset` based on access patterns. `MapDataset` supports efficient random access and is suitable for debugging or when order-dependent operations are needed. `IterDataset` (often created via `MapDataset.to_iter_dataset()`) is designed for performant, sequential iteration, typically used for training loops, especially with prefetching.fixFor random access or debugging, use `grain.MapDataset`. For performance-critical iteration during training, convert to `grain.IterDataset` using `dataset.to_iter_dataset()`.
affects: All versions
Errors
Common errors & fixes
ModuleNotFoundError: No module named 'grain.tensorflow'
The `grain.tensorflow` module or its functionality has likely been moved, refactored, or removed in `grain` version 0.2.16 or recent updates, or it's not part of the standard `grain` distribution, which primarily emphasizes JAX.
fixCheck the official Grain documentation or changelog for the correct way to integrate with TensorFlow, or if the functionality has been replaced or moved to a different module. It's possible the TensorFlow integration approach has changed.
ModuleNotFoundError: No module named 'grain.python.experimental'
Experimental APIs previously located under `grain.python.experimental` have been deprecated and their functionality has been moved to more stable, top-level `grain.IterDataset` or `grain.MapDataset` methods in `grain` 0.2.16.
fixMigrate usage from `grain.python.experimental.MultiprocessPrefetchIterDataset` to `grain.IterDataset.mp_prefetch` and from `grain.python.experimental.ConcatenateMapDataset` to `grain.MapDataset.concatenate`.
TypeError: __getitem__() missing 1 required positional argument: 'index'
In `grain` version 0.2.16, custom implementations of `RandomAccessDataSource.__getitem__` are now expected to explicitly accept an `int` index. Passing an object of type `typing.SupportsIndex` instead of `int` can cause this runtime error or type checking issues.
fixUpdate the `__getitem__` method signature in custom `RandomAccessDataSource` implementations to explicitly accept an `int` argument: `def __getitem__(self, index: int):`.
ValueError: Cannot batch a filtered MapDataset. Convert to IterDataset first using .to_iter_dataset().
`MapDataset.batch` cannot directly follow `MapDataset.filter` because the `filter` transformation can introduce `None` elements (representing removed items). The batching operation on a `MapDataset` expects contiguous, non-None elements. Converting to `IterDataset` after filtering ensures `None` elements are properly skipped before batching.
fixInsert `.to_iter_dataset()` between the `.filter()` and `.batch()` calls: `ds = ds.filter(lambda x: x % 2 == 0).to_iter_dataset().batch(2)`.
ERROR: No matching distribution found for grain
`grain` version 0.2.16 (and newer) requires Python >=3.11. This error occurs when attempting to install the library using `pip` with an unsupported Python version (e.g., Python 3.10) because no compatible wheel exists for that Python version.
fixUpgrade your Python environment to version 3.11 or newer. Use a Python environment manager (like `conda` or `pyenv`) to create or switch to a compatible Python version.
Upgrade
Version history
0.2.18latest on PyPI · released Jun 17, 2026
Audit
Dependencies
numpyrequiredCommon dependency for data manipulation.
absl-pyrequiredGoogle's Python Abseil library, often used in Google projects.
jaxoptionalPrimary target framework, though Grain does not strictly require JAX to run.
array-recordoptionalFor reading data from ArrayRecord format.
pyarrowoptionalFor reading data from Parquet files via ParquetIterDataset.
tensorflow-datasetsoptionalFor integrating with TensorFlow Datasets.
orbax-checkpointoptionalFor asynchronous checkpointing of data loading state.