Skip to content

Roxxel

Welcome to the documentation for Roxxel: a zero-RAM, JAX-centric dataloading, streaming, and checkpointing & logging toolkit.

Roxxel is designed specifically for large-scale deep learning pipelines using JAX & Flax NNX. It provides OS-level memory mapping, zero-copy sharded device streaming, background model checkpointing, and thread-safe distributed logging.

📦 Installation

Roxxel can be installed via pip directly from PyPI.

To install the core dataloader and async logging engine only: 

pip install roxxel

To install the JAX-native asynchronous checkpointing extensions: 

pip install roxxel[checkpoint]

  • Tutorial / Cookbook: A complete end-to-end JAX/Flax NNX training harness.
  • Dataloader & Streaming: Zero-RAM block loading and JAX-native multi-device data placement.
  • Checkpointing: Non-blocking weight and optimizer state preservation using Orbax.
  • Logging: Zero-overhead logging with multi-host safety and automatic traceback capture.