cyreal 0.2.0

Jittable data loading utilities for JAX.

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))
iterate = jax.jit(loader.next)

for epoch in range(2):
    for _ in range(loader.steps_per_epoch):
        batch, state, mask = iterate(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

Library	Dataset Device	Batch Device	Method	Time (s)
grain	CPU	CPU	Iterator	1.33
cyreal	CPU	CPU	jit(loader.next)	0.04
cyreal	CPU	CPU	scan_epoch	0.09

A40 with Wimpy CPU

Library	Dataset Device	Batch Device	Method	Time (s)
grain	CPU	CPU	Iterator	10.34
grain	CPU	GPU	Iterator	11.65
cyreal	CPU	CPU	jit(loader.next)	0.66
cyreal	CPU	GPU	jit(loader.next)	0.68
cyreal	GPU	GPU	jit(loader.next)	0.66
cyreal	CPU	CPU	scan_epoch	3.78
cyreal	CPU	GPU	scan_epoch	4.00
cyreal	GPU	GPU	scan_epoch	4.35

RTX 5090

Library	Dataset Device	Batch Device	Method	Time (s)
grain	CPU	CPU	Iterator	3.80
grain	CPU	GPU	Iterator	4.04
cyreal	CPU	CPU	jit(loader.next)	0.50
cyreal	CPU	GPU	jit(loader.next)	0.50
cyreal	GPU	GPU	jit(loader.next)	0.50
cyreal	CPU	CPU	scan_epoch	2.71
cyreal	CPU	GPU	scan_epoch	2.72
cyreal	GPU	GPU	scan_epoch	2.68