Data and batching

Construct datasets, reproduce shuffles, handle incomplete batches, and resume an iterator without changing the sample order.

These outputs were saved from local execution. The documentation renderer does not execute notebook cells.

Show code cell source

Hide code cell source

from pathlib import Path
import sys
ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent
sys.path.insert(0, str(ROOT))
import jax
import jax.numpy as jnp
import numpy as np
import jaxsaw as js
from jaxsaw import nn, optim, data, checkpoint, inspection
from jaxsaw.nn import functional as F
print(f"JAXsaw {js.__version__} · JAX {jax.__version__}")
print(jax.devices())

JAXsaw 0.1.0 · JAX 0.5.3
[CudaDevice(id=0)]

Construct a dataset

A TensorDataset returns tuples with matching leading dimensions. DataLoader batches those samples into JAX arrays. Keep normalization statistics fitted on the training split.

features = np.arange(30,dtype=np.float32).reshape(10,3)
targets = np.arange(10,dtype=np.int32)
dataset = data.TensorDataset(features,targets)
loader = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)
print('Examples:',len(dataset))
print('First sample:',dataset[0])

Examples: 10
First sample: (array([0., 1., 2.], dtype=float32), np.int32(0))

Inspect an epoch

The final batch is retained by default. Different shapes may compile separately; use drop_last for training or padding and a valid-item mask when appropriate. Never silently discard evaluation examples.

ids = []
for x,y in loader.iter_epoch(0):
    print('Batch:',x.shape,'sample IDs:',np.asarray(y))
    ids.extend(np.asarray(y).tolist())
assert sorted(ids)==list(range(10))

Batch: (4, 3) sample IDs: [7 1 5 0]
Batch: (4, 3) sample IDs: [3 4 2 6]
Batch: (2, 3) sample IDs: [9 8]

Reproduce shuffling

Shuffle order is derived from the loader seed and epoch. Identical configuration gives identical order; a new epoch changes it.

def order(epoch):
    copy = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)
    return np.concatenate([np.asarray(y) for _,y in copy.iter_epoch(epoch)])
print('Epoch 0:',order(0))
print('Epoch 1:',order(1))
assert np.array_equal(order(0),np.asarray(ids))

Epoch 0: [7 1 5 0 3 4 2 6 9 8]
Epoch 1: [1 0 8 3 7 2 5 4 6 9]

Resume after a completed update

The loader cursor advances before yielding. Save it only after the corresponding optimizer update completes. A restored iterator continues with the next batch.

loader = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)
iterator = loader.iter_epoch(0)
first = next(iterator)
state = loader.state_dict()
expected = next(iterator)[1]
restored = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)
restored.load_state_dict(state)
actual = next(restored.iter_epoch(0))[1]
print('Restored next IDs:',np.asarray(actual))
assert np.array_equal(expected,actual)

Restored next IDs: [3 4 2 6]

Pad variable-length sequences

Return explicit lengths and masks from a custom collator. Recurrent layers use lengths to freeze the carry; token losses use a valid-item mask to normalize correctly.

def pad_sequences(samples):
    sequences, labels = zip(*samples)
    lengths = np.array([len(s) for s in sequences],dtype=np.int32)
    padded = np.zeros((len(samples),int(lengths.max()),1),dtype=np.float32)
    for i,sequence in enumerate(sequences):
        padded[i,:len(sequence),0]=sequence
    return jnp.asarray(padded),jnp.asarray(lengths),jnp.asarray(labels)
batch = pad_sequences([(np.array([1.,2.,3.]),0),(np.array([4.]),1)])
print('Padded shape:',batch[0].shape,'lengths:',batch[1])
recurrent = nn.LSTM(1,4,rngs=js.Rngs(params=0))
sequence,carry = recurrent(batch[0],lengths=batch[1])
print('Sequence:',sequence.shape,'hidden/cell:',[a.shape for a in carry])

Padded shape: (2, 3, 1) lengths: [3 1]
Sequence: (2, 3, 4) hidden/cell: [(1, 2, 4), (1, 2, 4)]