Grain - Feeding JAX Models Grain is a library for reading data for training and evaluating JAX models. It's open source, fast and deterministic. Installation: pip install grain Docs Grain is used by MaxText, a simple, performant and scalable JAX codebase for LLM.