Grain - Feeding JAX Models
Installation | Quickstart | Reference docs | Change logs
Grain is a Python library for reading and processing data for training and evaluating JAX models. It is flexible, fast and deterministic.
Grain allows to define data processing steps in a simple declarative way:
import grain
dataset = (
grain.MapDataset.source([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
.shuffle(seed=42) # Shuffles elements globally.
.map(lambda x: x+1) # Maps each element.
.batch(batch_size=2) # Batches consecutive elements.
)
for batch in dataset:
# Training step.
Grain is designed to work with JAX models but it does not require JAX to run and can be used with other frameworks as well.
Installation
Grain is available on PyPI and can be
installed with pip install grain.
Supported platforms
Grain does not directly use GPU or TPU in its transformations, the processing within Grain will be done on the CPU by default.
| Linux | Mac | Windows | |
|---|---|---|---|
| x86_64 | yes | no | yes |
| aarch64 | yes | yes | n/a |
Quickstart
Citing Grain
To cite this repository:
@software{grain2023github,
author = {Marvin Ritter and Ihor Indyk and Aayush Singh and Andrew Audibert and Anoosha Seelam and Camelia Hanes and Eric Lau and Jacek Olesiak and Jiyang Kang and Xihui Wu},
title = {{Grain} - Feeding JAX Models},
url = {http://github.com/google/grain},
version = {0.2.12},
year = {2023},
}
The version number is intended to be that from pyproject.toml, and the year corresponds to the project's open-source release.
Existing users
Grain is used by MaxText, Gemma, kauldron, maxdiffusion and multiple internal Google projects.
Release files for grain 0.2.19
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distributions (wheels)
Total release size: 9.0 MB
Release files / grain-0.2.19-cp314-cp314-win_amd64.whl
| Download URL | grain-0.2.19-cp314-cp314-win_amd64.whl |
|---|---|
| Size | 538.5 kB |
| Tags | CPython 3.14 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
dba1ea22d143275b26da24e26f34f2fec908956dba0ed2fbc40345aa8cdf0e8c
|
|
BLAKE2b-256 checksum How to use checksums |
51b2cd180d137a8b92a049fbfdc74c461b55fbd744f734a40be97f179f4d2e7e
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
| Download URL | grain-0.2.19-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl |
|---|---|
| Size | 592.9 kB |
| Tags | CPython 3.14 Linux glibc 2.24+ x86-64 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
2a67e1c41949feb669461e8de51c84d820b61259bd4aac21691227125ec9705e
|
|
BLAKE2b-256 checksum How to use checksums |
1e49119cf8c064b7aff6e97b1ad0863dfd035493373e069b7051fba4b273bcb3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
| Download URL | grain-0.2.19-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl |
|---|---|
| Size | 591.7 kB |
| Tags | CPython 3.14 Linux glibc 2.24+ ARM64 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
7007978b35f30c85105bd210489e0fe0b87fcd6e73c782afd20dac944526d0ee
|
|
BLAKE2b-256 checksum How to use checksums |
c2a1589cfe8077cf3e51e25a653b838ce47f583d99ad87d518958d078a2318d8
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp314-cp314-macosx_11_0_arm64.whl
| Download URL | grain-0.2.19-cp314-cp314-macosx_11_0_arm64.whl |
|---|---|
| Size | 535.3 kB |
| Tags | CPython 3.14 macOS 11.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
f5bc250f4e7a0f84a575c4977cfb7b38da4e837ed346939cb4a86f4cab12d2ea
|
|
BLAKE2b-256 checksum How to use checksums |
71f4d6aebdca618cef69ae04efd82353272e4832c5d69a5e045c283807d8cbe8
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp313-cp313-win_amd64.whl
| Download URL | grain-0.2.19-cp313-cp313-win_amd64.whl |
|---|---|
| Size | 532.3 kB |
| Tags | CPython 3.13 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
7beb231f20b462491cd27e09b8cd909a7ac4669b04cdec454137cfc14945b7f8
|
|
BLAKE2b-256 checksum How to use checksums |
0713b289dfe38c6a7ba270bd9e5fe9255f6863cc350ba900022b705d3647a655
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
| Download URL | grain-0.2.19-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl |
|---|---|
| Size | 592.8 kB |
| Tags | CPython 3.13 Linux glibc 2.24+ x86-64 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
febb51a2a0401e832c74d2765c2ae541893e61df809d13f92af77314deae7074
|
|
BLAKE2b-256 checksum How to use checksums |
ac029cf846b7c4c40160bf83117e46d74355128dfda489c4f1a4f83c81b54ed2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
| Download URL | grain-0.2.19-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl |
|---|---|
| Size | 591.6 kB |
| Tags | CPython 3.13 Linux glibc 2.24+ ARM64 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
8acbaa2bf481e08b17d1f500a7f820c4f58d22105c395ff3d2887d202eaa6261
|
|
BLAKE2b-256 checksum How to use checksums |
9de0eb5c108af11f1fc306add51f131109f147f46db3b4676674e057c7b47647
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp313-cp313-macosx_11_0_arm64.whl
| Download URL | grain-0.2.19-cp313-cp313-macosx_11_0_arm64.whl |
|---|---|
| Size | 535.3 kB |
| Tags | CPython 3.13 macOS 11.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
6fef969763f2ce8e90a5c39d83ec26ee365ec99eb2bb0a1c17fea725d45467e3
|
|
BLAKE2b-256 checksum How to use checksums |
00a2d342f70f9dfada2d61da48973466880d835c3cb16465f5f271cf24f09452
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp312-cp312-win_amd64.whl
| Download URL | grain-0.2.19-cp312-cp312-win_amd64.whl |
|---|---|
| Size | 532.2 kB |
| Tags | CPython 3.12 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
3cacb666b5dfaea24b7be7ba74e0dee1e34651029a0f4b3890233a3f405855a0
|
|
BLAKE2b-256 checksum How to use checksums |
721a5a15c84ab06d99c9f7009b87a59727389e42e88d280d6e00486f81bda38a
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
| Download URL | grain-0.2.19-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl |
|---|---|
| Size | 592.7 kB |
| Tags | CPython 3.12 Linux glibc 2.24+ x86-64 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
b3a17f4c0484c487316ccec867cba4b0870d15bfd7bd2de8c1cd4f03afd8b570
|
|
BLAKE2b-256 checksum How to use checksums |
96230b61cc3ef5046d99616823cd614304119e5756b944f59aada3a7c3bd5965
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
| Download URL | grain-0.2.19-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl |
|---|---|
| Size | 591.5 kB |
| Tags | CPython 3.12 Linux glibc 2.24+ ARM64 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
6db658e4e509fbfc4f4bbd3c671d14abfe7bf51a146dd64ac85798831dfca23a
|
|
BLAKE2b-256 checksum How to use checksums |
6541210fcdd7e2306c6b97f237a86da23a0a698da84d5a2681153a38cd02fad4
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp312-cp312-macosx_11_0_arm64.whl
| Download URL | grain-0.2.19-cp312-cp312-macosx_11_0_arm64.whl |
|---|---|
| Size | 535.1 kB |
| Tags | CPython 3.12 macOS 11.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
9395180b16a5d3653c628c6186b4a3e39d94d7e7bba6b62210afc3211447c21d
|
|
BLAKE2b-256 checksum How to use checksums |
aacfc4189db55da326a047c76c6aa947e61b6a0872982f20ab78ccb1d79f2f8c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp311-cp311-win_amd64.whl
| Download URL | grain-0.2.19-cp311-cp311-win_amd64.whl |
|---|---|
| Size | 532.0 kB |
| Tags | CPython 3.11 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
6d9956516ecb12a7c576ba8bd30f57572d9acf7cfcad3609a7dc4edc858f5c06
|
|
BLAKE2b-256 checksum How to use checksums |
4f36cd45784ced6eb8286fe560556434a612f2e7744d0fe824451c9569c03874
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
| Download URL | grain-0.2.19-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl |
|---|---|
| Size | 594.6 kB |
| Tags | CPython 3.11 Linux glibc 2.24+ x86-64 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
d268d3b943b368cd9c7c29bb9f168d78d61b4d18450c9f53374d6ad2a02bd505
|
|
BLAKE2b-256 checksum How to use checksums |
89b35812b040d8899c7eb8ea48b2abac45f27b78ae60c555c617e4b9843ac1eb
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
| Download URL | grain-0.2.19-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl |
|---|---|
| Size | 593.1 kB |
| Tags | CPython 3.11 Linux glibc 2.24+ ARM64 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
c9949d6c17f91aa03ce2dca199351380d2ec0ca199a09cab31dd26d4e67ba377
|
|
BLAKE2b-256 checksum How to use checksums |
c0dd21ae4b49bea768583a64f832c712977c49d855a5411534c10cdc3c8fc1e6
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Release files / grain-0.2.19-cp311-cp311-macosx_11_0_arm64.whl
| Download URL | grain-0.2.19-cp311-cp311-macosx_11_0_arm64.whl |
|---|---|
| Size | 537.7 kB |
| Tags | CPython 3.11 macOS 11.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
13c9efc019b72c524190913e1e555670268542f30213035d5744d042e782ed8d
|
|
BLAKE2b-256 checksum How to use checksums |
b1fb51753c5d6db1c787fe294301af138792241951f0c85a10871be8b0c9966d
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.16
|
Security for grain 0.2.19
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
