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.

Check out tutorials/ for more information on how to use Grain!

Name		Name	Last commit message	Last commit date
Latest commit History 385 Commits
.github/workflows		.github/workflows
docs		docs
grain		grain
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
.readthedocs.yaml		.readthedocs.yaml
BUILD		BUILD
CONTRIBUTING.md		CONTRIBUTING.md
LICENSE		LICENSE
README.md		README.md
WORKSPACE		WORKSPACE
pyproject.toml		pyproject.toml
setup.py		setup.py

Provide feedback