Grain: A library for loading and transforming data for ML training.

Navigation

Verified details

These details have been verified by PyPI
Project links
GitHub Statistics
Maintainers
Avatar for aireenmei from gravatar.com aireenmei Avatar for aseelam from gravatar.com aseelam Avatar for iindyk from gravatar.com iindyk

Unverified details

These details have not been verified by PyPI
Meta
  • License: Apache License
  • Author: Grain team
  • Requires: Python >=3.11
  • Provides-Extra: parquet
Classifiers
Report project as malware

Project description

Grain - Feeding JAX Models

Continuous integration PyPI version

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.

Project details

Verified details

These details have been verified by PyPI
Project links
GitHub Statistics
Maintainers
Avatar for aireenmei from gravatar.com aireenmei Avatar for aseelam from gravatar.com aseelam Avatar for iindyk from gravatar.com iindyk

Unverified details

These details have not been verified by PyPI
Meta
  • License: Apache License
  • Author: Grain team
  • Requires: Python >=3.11
  • Provides-Extra: parquet
Classifiers

Release history Release notifications | RSS feed

This version

0.2.16

0.2.15

0.2.14

0.2.13

0.2.12

0.2.11

0.2.10

0.2.10b1 pre-release

0.2.9

0.2.8

0.2.7

0.2.6

0.2.5 yanked

Reason this release was yanked:

mp_prefetch is broken.

0.2.3

0.2.2

0.2.1

0.2.0

0.1.0

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

grain-0.2.16-cp314-cp314-win_amd64.whl (520.6 kB view details)

Uploaded CPython 3.14Windows x86-64

grain-0.2.16-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (583.0 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

grain-0.2.16-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (583.6 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

grain-0.2.16-cp314-cp314-macosx_11_0_arm64.whl (526.9 kB view details)

Uploaded CPython 3.14macOS 11.0+ ARM64

grain-0.2.16-cp313-cp313-win_amd64.whl (512.7 kB view details)

Uploaded CPython 3.13Windows x86-64

grain-0.2.16-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (582.9 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

grain-0.2.16-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (583.6 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

grain-0.2.16-cp313-cp313-macosx_11_0_arm64.whl (526.9 kB view details)

Uploaded CPython 3.13macOS 11.0+ ARM64

grain-0.2.16-cp312-cp312-win_amd64.whl (512.7 kB view details)

Uploaded CPython 3.12Windows x86-64

grain-0.2.16-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (582.9 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

grain-0.2.16-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (583.6 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

grain-0.2.16-cp312-cp312-macosx_11_0_arm64.whl (526.8 kB view details)

Uploaded CPython 3.12macOS 11.0+ ARM64

grain-0.2.16-cp311-cp311-win_amd64.whl (512.3 kB view details)

Uploaded CPython 3.11Windows x86-64

grain-0.2.16-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (584.2 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

grain-0.2.16-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (585.3 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

grain-0.2.16-cp311-cp311-macosx_11_0_arm64.whl (528.8 kB view details)

Uploaded CPython 3.11macOS 11.0+ ARM64

File details

Details for the file grain-0.2.16-cp314-cp314-win_amd64.whl.

File metadata

  • Download URL: grain-0.2.16-cp314-cp314-win_amd64.whl
  • Upload date:
  • Size: 520.6 kB
  • Tags: CPython 3.14, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for grain-0.2.16-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 143955414ceb8bc8df02599d83d006b5dfbc9eef55fd1dd729eb2ddefa14f821
MD5 9dc81abdd88e682aa437ae6b99780123
BLAKE2b-256 c2468a89a61150a298e001eb76f4825cff66324692683beacc47a8efefc26cf8

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 aa19d0ddcca020c7f78bfdaa480304e943f363061be273af216c7334fe504269
MD5 b73fc42988dd041efaec2ad7672b232e
BLAKE2b-256 a88ebe21ee798cba3c7568198a97c8e8ebe07d8074e504c06088ec4358a7c949

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a70df579693ac1068ab646ec770ac1f4fd1456b06ed79089c5aa3b2d9a0ca6c2
MD5 5cb5f73431b08284a79f88fee2fa2a07
BLAKE2b-256 a0ddd9766f03b8abc8ea8bca46b1c0deb4ce8fff7aadb0e2cae60bc1daae7c0d

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp314-cp314-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp314-cp314-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 59262b5d74a3e69c209ffe2f1f599db9063ac9114f19ca93d43a17939d3e96dc
MD5 fa97c5a761a32d689bb196d0d853bcc2
BLAKE2b-256 394c24ead40ebc98184189895a3b71910e50988a6b8af2730d49df6f98b74101

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp313-cp313-win_amd64.whl.

File metadata

  • Download URL: grain-0.2.16-cp313-cp313-win_amd64.whl
  • Upload date:
  • Size: 512.7 kB
  • Tags: CPython 3.13, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for grain-0.2.16-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 c12a8d851cb380e5ca4f639acfade1a781f2c33409eb9cf876c23b3c25c7e7a4
MD5 d27996e5725e7193e433f99f1f138941
BLAKE2b-256 dc971df4ba25d711c83f7ae44d1b76206196dd42a16106229f66717d7cb71655

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 33440bc6e2253650b915c9846ba541ac382169280c175fa5e464664c13c0f16f
MD5 492aa42658fbdef914efef65da1089bd
BLAKE2b-256 b6ed9bd6158f78e000b522a2da264b544b4cde24ece1b6e63a2a7699ed9413be

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 71876c72a2ed7ac0074561806e133725a0d28699ae259945546c05fb2d86de49
MD5 a3360fac6c76b0142917e6c5ba0f10e7
BLAKE2b-256 7ffdc9cc44225c8e2156e4191e1038c8be8dee9fe5edef184a592e2617bbafef

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp313-cp313-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp313-cp313-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 603a0e3e015690384395a8c0243ada4d543820208577294796a8ba0fb6283a8c
MD5 724af97ad0888f04948e5f4784e802c8
BLAKE2b-256 fe110169341b9475ab2682db43b7be40fe21f6841603bed2fa5be944091877df

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp312-cp312-win_amd64.whl.

File metadata

  • Download URL: grain-0.2.16-cp312-cp312-win_amd64.whl
  • Upload date:
  • Size: 512.7 kB
  • Tags: CPython 3.12, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for grain-0.2.16-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 b5000a7a21b6cac7ecd58ac26aea341fef1e30c125c99d0fb3524f32b0a61c55
MD5 78ad69890d281fdb1f77a02ff889d5e8
BLAKE2b-256 6ac9b940dbf99565fd1078902102ac51d9957b14921c993cf457c58f3e5c35c2

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 0ab43a9127724f67f37425b598974ce3c1dedb6d33ad3d3905dd84adcc5f9a95
MD5 e4fc7155edb8380be9579d699d2b65a9
BLAKE2b-256 6ad513209f236c0fec2bbcb32bd5110d7eb4faf15192b5b1e958bdf460e62ddb

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 f5d662e0fd4abce92e3fe31f2396e460890326b51a5452a74a46240c0874f05e
MD5 10ec99caf244fdf230906d64c955cf39
BLAKE2b-256 1908f36a8b062f99d8c340fce219e6864edae60446fce3c5b309ec8305653579

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp312-cp312-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 cd1817d3bc2bf66793167d0056c023d24a414dad2281acfc4625b73390ef4b9c
MD5 67a34a83130a1c1bb08509b989fda141
BLAKE2b-256 d6cc61e4305e89227416317b91053712cdac1418b0f5a482a7ccb319e21c46c4

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp311-cp311-win_amd64.whl.

File metadata

  • Download URL: grain-0.2.16-cp311-cp311-win_amd64.whl
  • Upload date:
  • Size: 512.3 kB
  • Tags: CPython 3.11, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for grain-0.2.16-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 3bed568195b225014a542a073030a327378283e327106cb58696f376ef9fd229
MD5 8b3fbe3cee85de2721932eaea5d8cb2c
BLAKE2b-256 9dc268611b5097cf43535320bc08503d989955a0815bbfb827d3c340c562d288

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 3d39b2a1f02da9ea67ccfae56a8720453629feb05f3e48cc6c32877804210a31
MD5 5af9d5224572e25a105d9bc4427e4297
BLAKE2b-256 7bf2f651b7e12d8f9c41c2b030972af1251f8608980409c4ce69e61f78b08458

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 6fd3d38317549e3714ee280cd4dee1373578b8c5c4e2e8ae4af4970714ab11ab
MD5 77426173bd8b67b763e90e07f6a12973
BLAKE2b-256 3241b13cdc14bd8c5eefe0f7436aa83674eae6bbc8c5faf686917b095238f518

See more details on using hashes here.

File details

Details for the file grain-0.2.16-cp311-cp311-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for grain-0.2.16-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 af572d514c8082f50cd4d05753e7b318a3539b11aa6cf2cd97914faf3fdff960
MD5 1bf072821144ff8802a5a84bc07cd90c
BLAKE2b-256 642e8aafa233d0c5d3a09a82406ca1fc9a20398dc8986039461b7ab91b73c303

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page