# Feeding the Beast: Data Pipelines with Grain

Over the last seven weeks we have [built models](https://kambale.dev/flax-nnx), written [training loops](https://kambale.dev/training-loop-in-jax), composed [optimizers](https://kambale.dev/optax-optimizers), bulletproofed our code with [Chex](https://kambale.dev/chex), and learned to debug inside [JIT-compiled functions](https://kambale.dev/debugging-jax). Every one of those articles used dummy data or a small built-in dataset, generated on the fly with a few lines of NumPy or `tfds.load()`.

Real projects do not work that way. Real datasets are millions of images sitting in cloud storage, gigabytes of audio clips streamed from disk, or CSV files too large to load into memory at once. And here is the uncomfortable truth: it does not matter how fast your `@nnx.jit` training step is if your data loader cannot keep up with it.

This is one of the most common and least glamorous bottlenecks in machine learning engineering. You buy an expensive GPU or rent TPU time, write a beautifully optimized training loop, and then watch your utilization graph sit at 20 percent because the accelerator is idle, waiting for the next batch of data to arrive from a slow Python loop.

Today we fix that.

We are introducing **Grain**, a data loading library purpose built for JAX. Grain plays a role similar to `torch.utils.data.DataLoader` in PyTorch, but its design accounts for the things JAX cares about most: determinism, explicit randomness, and staying out of the way of your compiled training step.

By the end of this article you will understand:

1.  Why naive Python data loading becomes the bottleneck in JAX training
    
2.  The three building blocks of a Grain pipeline: DataSource, Sampler, and Transformations
    
3.  How to write deterministic and reproducible data augmentation
    
4.  How to parallelize data loading across multiple CPU workers
    
5.  How to shard data correctly for distributed training
    
6.  How to wire a Grain pipeline into the training loop we built in Week 4
    

Let's feed the beast.

## Why data loading becomes the bottleneck

Picture a training loop where the model itself runs in under a millisecond per batch, thanks to `@nnx.jit` and XLA compilation. Now picture a plain Python function that reads an image from disk, decodes it, resizes it, and normalizes it. That single operation, done in pure Python, can easily take several milliseconds.

If your data loading is a simple `for` loop over a list of file paths, your training loop looks like this:

```python
for i in range(num_steps):
    # This line might take 5-10ms per image, done sequentially
    batch = load_and_preprocess_batch(file_paths[i])
    
    # This line might take 0.5ms thanks to JIT
    loss = train_step(model, optimizer, batch)
```

The GPU finishes its work almost instantly and then sits idle while the CPU slowly prepares the next batch. You paid for a race car and you are driving it in first gear.

The fix is not to make the model slower to match the data loader. The fix is to make data loading fast, parallel, and asynchronous, so that while the GPU is busy computing batch N, the CPU is already preparing batch N+1 in the background.

That is exactly what Grain is built to do.

## The three building blocks of Grain

A Grain pipeline is assembled from three components that work together: a **DataSource** that knows how to read individual records, a **Sampler** that decides the order records are visited in, and a list of **Transformations** that process each record or batch.

```python
import grain.python as grain

data_loader = grain.DataLoader(
    data_source=my_source,
    sampler=my_sampler,
    operations=my_transformations,
    worker_count=4,
)
```

Let's build each piece.

### The DataSource

A `DataSource` is the simplest possible abstraction: something with a length and a way to fetch a single item by index. It does not know anything about batching, shuffling, or augmentation. Its only job is: given an index, return the raw record.

```python
import grain.python as grain

class MyImageSource(grain.RandomAccessDataSource):
    def __init__(self, file_paths, labels):
        self._file_paths = file_paths
        self._labels = labels
    
    def __len__(self):
        return len(self._file_paths)
    
    def __getitem__(self, index):
        # Return raw, unprocessed data. Keep this fast and simple.
        return {
            'path': self._file_paths[index],
            'label': self._labels[index],
        }

source = MyImageSource(file_paths=my_paths, labels=my_labels)
```

Notice how minimal this is. It does not decode images or apply any transformations. That separation of concerns is intentional and it is what allows Grain to parallelize the expensive work later.

### The Sampler

The `Sampler` decides which index to visit, in what order, for how many epochs, and with what random seed. This is where shuffling and reproducibility live.

```python
sampler = grain.IndexSampler(
    num_records=len(source),
    shuffle=True,
    num_epochs=None,   # None means infinite, useful for continuous training
    seed=42,
)
```

Because the seed is explicit, two runs with the same seed visit records in exactly the same order. This matters enormously when you are debugging a training run or trying to reproduce a result from a paper or a colleague's experiment.

### Transformations

Transformations are the actual processing steps applied to each record: decoding an image, normalizing pixel values, applying augmentation, or grouping records into batches.

There are two kinds of custom transforms in Grain, and the distinction matters.

**Deterministic transforms** inherit from `MapTransform` and implement a `map` method. Given the same input, they always produce the same output.

```python
import numpy as np

class Normalize(grain.MapTransform):
    def map(self, element):
        element['image'] = element['image'].astype(np.float32) / 255.0
        return element
```

**Random transforms** inherit from `RandomMapTransform` and implement `random_map`, which receives a `numpy.random.Generator` as an argument. This is the critical detail: you must use the provided generator, not `np.random` directly, or your augmentation will not be reproducible.

```python
class RandomFlip(grain.RandomMapTransform):
    def random_map(self, element, rng):
        if rng.random() > 0.5:
            element['image'] = np.fliplr(element['image'])
        return element
```

You then assemble your transforms into an ordered list, finishing with `grain.Batch` to group individual records into batches:

```python
operations = [
    Normalize(),
    RandomFlip(),
    grain.Batch(batch_size=64, drop_remainder=True),
]
```

Notice `drop_remainder=True`. If you have been following along since Week 6, you will remember exactly why this matters: an inconsistent final batch size forces JAX to recompile your JIT function. Dropping the incomplete final batch keeps every shape identical and avoids the recompilation trap entirely.

## Putting it together

Here is a complete pipeline combining all three pieces:

```python
import grain.python as grain
import numpy as np

# DataSource: raw access to records
class MNISTSource(grain.RandomAccessDataSource):
    def __init__(self, images, labels):
        self._images = images
        self._labels = labels
    
    def __len__(self):
        return len(self._images)
    
    def __getitem__(self, index):
        return {
            'image': self._images[index],
            'label': self._labels[index],
        }

source = MNISTSource(images=train_images, labels=train_labels)

# Sampler: order and reproducibility
sampler = grain.IndexSampler(
    num_records=len(source),
    shuffle=True,
    num_epochs=None,
    seed=42,
)

# Transformations: processing pipeline
class Normalize(grain.MapTransform):
    def map(self, element):
        element['image'] = element['image'].astype(np.float32) / 255.0
        return element

operations = [
    Normalize(),
    grain.Batch(batch_size=64, drop_remainder=True),
]

# Assemble the DataLoader
data_loader = grain.DataLoader(
    data_source=source,
    sampler=sampler,
    operations=operations,
    worker_count=4,
)

# Use it as an iterator
data_iterator = iter(data_loader)
first_batch = next(data_iterator)

print(f"Batch image shape: {first_batch['image'].shape}")
print(f"Batch label shape: {first_batch['label'].shape}")
```

This looks similar to a PyTorch `DataLoader`, and that is intentional. The difference is under the hood: Grain's sampler and random transforms are built around the same explicit, seedable randomness philosophy that JAX itself uses.

## Parallelism: The worker\_count parameter

The single most impactful performance knob in Grain is `worker_count`.

```python
data_loader = grain.DataLoader(
    data_source=source,
    sampler=sampler,
    operations=operations,
    worker_count=0,   # Sequential, single process
)
```

With `worker_count=0`, everything runs sequentially in your main Python process. This is useful for debugging, since you can set breakpoints inside your transforms and step through them normally, but it will not use your CPU efficiently.

```python
data_loader = grain.DataLoader(
    data_source=source,
    sampler=sampler,
    operations=operations,
    worker_count=8,   # Parallel, using multiprocessing
)
```

With `worker_count > 0`, Grain spins up multiple worker processes. Because these are separate processes rather than threads, they bypass Python's Global Interpreter Lock entirely, meaning your image decoding, resizing, and augmentation can genuinely run in parallel across CPU cores. Grain uses shared memory to transfer completed batches back to your main process efficiently, avoiding the overhead of pickling large arrays repeatedly.

A practical rule of thumb: start with `worker_count` roughly equal to the number of CPU cores available, then benchmark. If your data source is already in memory (for example, a NumPy array you loaded once at the start), you may not need many workers at all, since there is no disk I/O to overlap. In that case:

```python
data_loader = grain.DataLoader(
    data_source=source,
    sampler=sampler,
    operations=operations,
    worker_count=0,
    read_options=grain.ReadOptions(num_threads=0),
)
```

Disabling thread prefetching when your dataset already lives in memory avoids unnecessary overhead from a prefetching mechanism you do not need.

## Sharding for distributed training

When you train across multiple JAX processes, whether that is multiple GPUs on one machine or multiple machines in a cluster, each process must see a distinct, non overlapping slice of the data. Otherwise you are wasting compute recomputing gradients on the same examples multiple times per step.

Grain handles this with `ShardOptions`, passed into your sampler:

```python
try:
    # Automatically detects jax.process_index() and jax.process_count()
    shard_options = grain.ShardByJaxProcess(drop_remainder=True)
except ImportError:
    # Fallback for a single process setup
    shard_options = grain.ShardOptions(
        shard_index=0, shard_count=1, drop_remainder=True
    )

sampler = grain.IndexSampler(
    num_records=len(source),
    shard_options=shard_options,
    shuffle=True,
    num_epochs=None,
    seed=42,
)
```

`ShardByJaxProcess` is the recommended default. It reads directly from the JAX runtime to figure out how many processes are running and which one this particular process is, then automatically hands each process its own disjoint slice of the dataset. You do not need to manually calculate offsets or worry about duplicated examples across workers.

We will return to this in much more depth in Week 10, when we cover sharding and parallelism for the model itself. For now, the important thing to remember is that data sharding and model sharding are two separate concerns, and Grain handles the data side cleanly.

## Rebuilding a training loop with Grain

Let's take the training loop from [week 4](https://kambale.dev/training-loop-in-jax) and swap the dummy NumPy data generation for a real Grain pipeline.

### The data pipeline

```python
import grain.python as grain
import numpy as np
import jax
import jax.numpy as jnp
from flax import nnx
import optax

class ArrayDataSource(grain.RandomAccessDataSource):
    """Wraps in-memory NumPy arrays as a Grain data source."""
    def __init__(self, images, labels):
        self._images = images
        self._labels = labels
    
    def __len__(self):
        return len(self._images)
    
    def __getitem__(self, index):
        return {
            'image': self._images[index],
            'label': self._labels[index],
        }


class Normalize(grain.MapTransform):
    def map(self, element):
        element['image'] = element['image'].astype(np.float32) / 255.0
        return element


def build_data_loader(images, labels, batch_size, seed):
    source = ArrayDataSource(images, labels)
    
    sampler = grain.IndexSampler(
        num_records=len(source),
        shuffle=True,
        num_epochs=None,
        seed=seed,
    )
    
    operations = [
        Normalize(),
        grain.Batch(batch_size=batch_size, drop_remainder=True),
    ]
    
    return grain.DataLoader(
        data_source=source,
        sampler=sampler,
        operations=operations,
        worker_count=4,
    )
```

### Model, optimizer, and training step

This should look familiar from week 4:

```python
class ArchiveDigitizer(nnx.Module):
    def __init__(self, hidden_dim: int, num_classes: int, rngs: nnx.Rngs):
        self.linear1 = nnx.Linear(784, hidden_dim, rngs=rngs)
        self.linear2 = nnx.Linear(hidden_dim, num_classes, rngs=rngs)
    
    def __call__(self, x):
        x = self.linear1(x)
        x = nnx.relu(x)
        return self.linear2(x)


rngs = nnx.Rngs(jax.random.PRNGKey(42))
model = ArchiveDigitizer(hidden_dim=128, num_classes=10, rngs=rngs)

tx = optax.adam(learning_rate=0.005)
optimizer = nnx.Optimizer(model, tx=tx, wrt=nnx.Param)


@nnx.jit
def train_step(model, optimizer, batch_images, batch_labels):
    def loss_fn(m):
        logits = m(batch_images)
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits=logits, labels=batch_labels
        ).mean()
        return loss
    
    loss_val, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss_val
```

### The training loop, powered by Grain

```python
# Assume train_images, train_labels are loaded NumPy arrays, e.g. from MNIST
batch_size = 64
data_loader = build_data_loader(train_images, train_labels, batch_size, seed=42)

data_iterator = iter(data_loader)
num_steps = 2000

print("Starting training with Grain data pipeline...")

running_loss = 0.0
for step in range(num_steps):
    batch = next(data_iterator)
    
    x_batch = jnp.array(batch['image']).reshape(batch_size, -1)
    y_batch = jnp.array(batch['label'])
    
    loss = train_step(model, optimizer, x_batch, y_batch)
    running_loss += loss.item()
    
    if step % 200 == 0 and step > 0:
        avg_loss = running_loss / 200
        print(f"Step {step:4d} | Average Loss: {avg_loss:.4f}")
        running_loss = 0.0

print("Training complete.")
```

The training step itself did not change at all from Week 4. Only the mechanism feeding it data changed, and that separation is exactly the point. Your model logic and your data pipeline should be independent, swappable pieces.

## Reproducibility and checkpointing

One benefit of Grain's explicit, seeded design is that data pipelines are fully reproducible. Given the same seed and the same source, two runs will visit records in the identical order and apply the identical augmentations at each step.

For long running jobs, you also want to be able to resume training after an interruption without skipping or repeating data. Grain integrates with Orbax, the checkpointing library we will cover in Week 9, allowing you to save the exact position of your data iterator alongside your model weights and optimizer state. When training resumes, both the model and the data pipeline pick up exactly where they left off.

## Common pitfalls

A few mistakes come up often enough to call out directly.

**Using np.random instead of the provided generator.** Inside a `RandomMapTransform`, always use the `rng` argument passed into `random_map`. Calling `np.random.random()` directly breaks reproducibility, since that global state is not controlled by Grain's seeding.

**Forgetting drop\_remainder.** As covered in Week 6, an inconsistent final batch shape triggers expensive JIT recompilation. Set `drop_remainder=True` on `grain.Batch` unless you have a specific reason not to.

**Overestimating worker\_count.** More workers is not always better. If your `__getitem__` implementation is already fast, for example reading from an in-memory array, spinning up eight worker processes adds multiprocessing overhead without any benefit. Benchmark before assuming.

**Doing heavy work in the DataSource instead of Transformations.** Keep `__getitem__` minimal. Expensive operations like image decoding and augmentation belong in your transformation list, where Grain can distribute them across workers. A bloated `__getitem__` cannot be parallelized as effectively.

## Exercises

1.  **Build a custom augmentation.** Write a `RandomMapTransform` that randomly rotates an image by a small angle, using the provided `rng` generator. Verify that running the pipeline twice with the same seed produces identical rotations.
    
2.  **Benchmark worker counts.** Run the same data pipeline with `worker_count` set to 0, 2, 4, and 8. Time how long it takes to iterate through 500 batches at each setting. Plot the results.
    
3.  **Simulate a distributed shard.** Manually construct two `ShardOptions` objects representing shard 0 of 2 and shard 1 of 2. Confirm that the two resulting samplers produce completely disjoint sets of indices across a full epoch.
    
4.  **Trigger and fix a recompilation.** Remove `drop_remainder=True` from your `grain.Batch` call. Train for a few epochs and use `chex.assert_max_traces` from Week 6 to confirm that the final, undersized batch triggers a second compilation. Then fix it.
    

## Quick reference

```python
import grain.python as grain
import numpy as np

# DataSource
class MySource(grain.RandomAccessDataSource):
    def __len__(self):
        return len(self._data)
    def __getitem__(self, index):
        return self._data[index]

# Sampler
sampler = grain.IndexSampler(
    num_records=len(source),
    shuffle=True,
    num_epochs=None,
    seed=42,
    shard_options=grain.ShardByJaxProcess(drop_remainder=True),
)

# Deterministic transform
class MyMap(grain.MapTransform):
    def map(self, element):
        return element

# Random transform
class MyRandomMap(grain.RandomMapTransform):
    def random_map(self, element, rng):
        return element

# Assembling the pipeline
data_loader = grain.DataLoader(
    data_source=source,
    sampler=sampler,
    operations=[MyMap(), MyRandomMap(), grain.Batch(64, drop_remainder=True)],
    worker_count=4,
)

data_iterator = iter(data_loader)
batch = next(data_iterator)
```

## What's next

Our training loop can now pull in real data at real speed. But every training run we have written so far lives and dies with the Python process running it. If your machine crashes, if you lose power, or if you simply need to stop and resume tomorrow, all of that progress vanishes.

Next week we solve that problem with **Orbax**, JAX's checkpointing library. We will learn how to save not just model weights, but the entire training state: optimizer momentum, step count, and even the data iterator position from Grain, so a training run can be paused and resumed without losing a single batch of progress.

**Next week**: *Never lose your progress: Checkpointing with Orbax*