Skip to content

google/grain

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

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.