Skip to content

Data Loading

dataloader

BaseDataLoader dataclass 

BaseDataLoader(
    X: NpJaxArray,
    y: NpJaxArray,
    batch_size: int = 64,
    to_jax: bool = True,
)

Bases: ABC

Abstract base class for NumPy/JAX DataLoaders.

Defines the common interface for all data loading implementations.

Parameters:

  • X (NpJaxArray) –

    Input data of shape (N, ...).

  • y (NpJaxArray) –

    Target data of shape (N, ...). Must match X in first dimension.

  • batch_size (int, default: 64 ) –

    Number of samples per batch.

  • to_jax (bool, default: True ) –

    If True, convert returned batches to JAX arrays.

reset abstractmethod 

reset() -> None

Reset the iterator state.

DataLoader dataclass 

DataLoader(
    X: NpJaxArray,
    y: NpJaxArray,
    batch_size: int = 64,
    to_jax: bool = True,
    shuffle: bool = False,
    drop_last: bool = False,
    seed: int | None = None,
)

Bases: BaseDataLoader

Simple NumPy/JAX DataLoader with optional shuffling.

This loader provides standard mini-batch iteration for tabular or tensor data. It supports deterministic iteration order and optional shuffling at every reset/epoch.

Returned batches have the structure:

X shape: (batch_size, *X.shape[1:])
y shape: (batch_size, *y.shape[1:])

Parameters:

  • X (NpJaxArray) –

    Input data of shape (N, ...).

  • y (NpJaxArray) –

    Target data of shape (N, ...). Must match X in first dimension.

  • batch_size (int, default: 64 ) –

    Number of samples per batch.

  • to_jax (bool, default: True ) –

    If True, convert returned batches to JAX arrays.

  • shuffle (bool, default: False ) –

    If True, reshuffle sample order on every reset.

  • drop_last (bool, default: False ) –

    If drop_last=False, the final batch may be smaller than batch_size when N is not divisible by batch_size. If drop_last=True, this incomplete final batch is skipped.

  • seed (int | None, default: None ) –

    RNG seed used for reproducible shuffling.

Raises:

  • ValueError –

    If batch_size <= 0, inputs are empty, or X and y have mismatched leading dimensions.

Examples:

>>> import numpy as np
>>> X = np.arange(10).reshape(10, 1)
>>> y = np.arange(10, 20).reshape(10, 1)
>>> loader = DataLoader(X, y, batch_size=4, shuffle=False)
>>> len(loader)
3
>>> xb, yb = next(iter(loader))
>>> xb.shape
(4, 1)
>>> yb.shape
(4, 1)

reset 

reset() -> None

Reset the iterator state without re-instantiating.

TSDataLoader dataclass 

TSDataLoader(
    X: NpJaxArray,
    y: NpJaxArray,
    batch_size: int = 64,
    to_jax: bool = True,
    seq_len: int = 1,
    pad_value: float = 0.0,
    drop_last: bool = False,
)

Bases: BaseDataLoader

Time-Series DataLoader for sequential data with historical context.

This loader automatically prepends historical context to each batch, enabling models to access previous timesteps. The returned batches have the structure:

X shape: (batch_size + seq_len, *X.shape[1:])  - includes history
y shape: (batch_size, *y.shape[1:])            - targets for the batch

The first seq_len rows of each X batch represent historical context that precedes the actual batch samples, allowing temporal models to access prior information.

Parameters:

  • X (NpJaxArray) –

    Input data of shape (N, ...). Sequential data samples.

  • y (NpJaxArray) –

    Target data of shape (N, ...). Must match X in first dimension.

  • batch_size (int, default: 64 ) –

    Number of samples per batch (excluding historical context).

  • to_jax (bool, default: True ) –

    If True, convert batches to JAX arrays. If False, keep as NumPy arrays.

  • seq_len (int, default: 1 ) –

    Number of historical timesteps to prepend to each batch. seq_len=1 means 1 row of history before the batch.

  • pad_value (float, default: 0.0 ) –

    Value used for padding at the beginning (no prior history available).

  • drop_last (bool, default: False ) –

    If drop_last=False, the final batch may be smaller than batch_size when N is not divisible by batch_size. If drop_last=True, this incomplete final batch is skipped.

Raises:

  • ValueError –

    If batch_size <= 0, X and y have different lengths, data is empty, or seq_len is negative.

Examples:

>>> import numpy as np
>>> X = np.arange(20).reshape(20, 1)
>>> y = np.arange(20, 40).reshape(20, 1)
>>> loader = TSDataLoader(X, y, batch_size=4, seq_len=2, drop_last=False)
>>> len(loader)
5
>>> xb, yb = next(iter(loader))
>>> xb.shape  # (batch_size + seq_len, 1) = (4 + 2, 1)
(6, 1)
>>> yb.shape  # (batch_size, 1) = (4, 1)
(4, 1)

reset 

reset() -> None

Reset the iterator state without re-instantiating.