grain

Grain: A library for loading and transforming data for neural network training.


License
Other
Install
pip install grain==0.1.1

Documentation

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-nightly
  • Docs
  • Grain is used by MaxText, a simple, performant and scalable JAX codebase for LLM.

This is not an officially supported Google product.