Welcome to JAX DataLoader's documentation!

JAX DataLoader is a high-performance data loading library for JAX that provides efficient data loading, preprocessing, and augmentation capabilities.

Installation

pip install jax-dataloaders

Quick Start

from jax_dataloader import DataLoader, DataLoaderConfig

# Create a configuration
config = DataLoaderConfig(
    batch_size=32,
    shuffle=True,
    num_workers=4
)

# Create a data loader
loader = DataLoader(config)

# Iterate over batches
for batch in loader:
    # Process your batch
    pass

Features

  • Efficient data loading with JAX

  • Support for various data formats (CSV, JSON, Images)

  • Data augmentation and preprocessing

  • Progress tracking and monitoring

  • Memory management and caching

  • Multi-GPU support

  • Memory monitoring and optimization

Contents

Indices and tables