Google's data pipeline for JAX that doesn't need JAX

What it does

Grain chains data transformations—shuffle, map, batch—through a declarative MapDataset API. It handles ingestion and preprocessing for ML training, keeping all transformation work on the CPU while accelerators handle the math. Google uses it to feed models like Gemma and MaxText.

The interesting bit

Despite the “Feeding JAX Models” branding, Grain deliberately avoids GPU and TPU dependencies in its transformations, so you can use it with PyTorch or TensorFlow without installing JAX. That flexibility is unusual for a library so tightly coupled to Google’s JAX ecosystem in production.

Key highlights

• Declarative API: chain .shuffle(), .map(), and .batch() on a MapDataset source
• Deterministic by design, which helps with reproducible training runs
• CPU-only transformations by default; won’t compete with your accelerator for memory
• Already adopted by Gemma, MaxText, kauldron, maxdiffusion, and internal Google projects
• Cross-platform for Linux and Windows (x86_64) and Linux/Mac (aarch64), though Mac x86_64 is explicitly unsupported

Caveats

• The README calls the library “fast” but provides no benchmarks or comparisons to back it up
• macOS on Intel (x86_64) is a hard “no” in the support matrix with no explanation given

Verdict

Try it if you want a deterministic preprocessor already vetted by Google’s own model teams. Pass if you need GPU-side data augmentation or a large community outside Google’s own repositories.