Building High-Performance Data Pipelines with Grain and ArrayRecord- Google Developers Blog
Oct 7, 2025Jiyang Kang, Shivaji Dutta, Ihor Indyk, Felix Chern

Summary
Introduces Grain, a flexible data loading library for JAX, and ArrayRecord, a highly efficient file format, to build high-performance data pipelines for large-scale machine learning. Grain ensures efficient data loading and preprocessing with multiprocessing and stateful, checkpointable iterators for reproducibility. ArrayRecord, built on Riegeli, offers efficient random access via a metadata index and parallel I/O through data chunks, enabling true global shuffling and significantly higher read throughput compared to sequential formats like TFRecord. The post also details methods for converting TFRecord datasets to ArrayRecord using `tfds build` or Apache Beam for scalable conversion.
More from the JAX Model Evaluation Metrics trail
- Accelerating Spatio-Temporal Attention for Video Diffusion on TPUs- Google Developers Blog
- Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs- Google Developers Blog
- Decoding cosmic signals with deep learning and Keras- Google Developers Blog
- Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU- Google Developers Blog
- HeyGen x Google Cloud: Bringing Avatar IV to TPUs- Google Developers Blog