Cyreal - Another JAX DataLoader

grain for the corporations, cyreal for the people

Pure jax utilities for iterating over finite and infinite datasets without ever touching torch or tensorflow. Dataloaders are fast and support jax.jit, jax.grad, jax.lax.scan, and other function transformations.

Installation

pip install cyreal

The only dependency is jax.

Quick Start

Write fast dataloaders without torch or tensorflow

import jax
import jax.numpy as jnp

from cyreal.transforms import BatchTransform, DevicePutTransform
from cyreal.loader import DataLoader
from cyreal.sources import ArraySource
from cyreal.datasets import MNISTDataset

train_data = MNISTDataset(split="train").as_array_dict()
pipeline = [
  # Load dataset into memory-backed array
  ArraySource(train_data, ordering="shuffle"),
  # Batch it
  BatchTransform(batch_size=128),
  # Move the batch to the GPU
  DevicePutTransform(),
]
loader = DataLoader(pipeline)
state = loader.init_state(jax.random.key(0))

for epoch in range(2):
    for _ in range(loader.steps_per_epoch):
        batch, state, mask = jax.jit(loader.next)(state)
        ... # Train your network

Use scan_epoch to jit and avoid boilerplate

model_state = {"params": jnp.array(0)}

def update(model_state, batch, mask):
    model_state = {"params": model_state['params'] + 1}
    return model_state, None

for epoch in range(2):
    state, model_state, _ = loader.scan_epoch(state, model_state, update)

Examples and Documentation

See our documentation for more examples.

• Do you enjoy premature optimization? Why not jit the entire training epoch?
• For the dirty and impure, we support logging metrics from within a jitted loader.
• Got yourself a huge dataset? Stream from a disk-backed source.
• Afraid of finite datasets? We provide gymnax-backed data sources for online reinforcement learning.
• Starving researcher/temporarily embarrassed hyperscaler? We support continual learning via reservoir sampling and replay buffers.

We also provide full end to end training examples

• MNIST
• Time Series
• Reinforcement Learning

Speed Test

You can compare the speed to the grain dataloader using this script. This is how long it takes to iterate though one epoch of MNIST

MacBook M4 Pro

A40 with Wimpy CPU

RTX 5090