Core Components

September 17, 2026 · View on GitHub

Core abstractions and building blocks that form the foundation of Datarax pipelines. These modules define the protocols, base classes, and data structures used throughout the framework.

Component Overview

Component	Purpose	Key Classes
Element & Batch	Data containers	`Element`, `Batch`, `Metadata`
Config	Typed configuration	`OperatorConfig`, `StructuralConfig`
Modules	Base abstractions	`DataraxModule`, `OperatorModule`
Protocols	Interface contracts	`DataSourceModule`, `SamplerModule`

!!! note "Key points"

- **Element** wraps a single data sample with state and metadata
- **Batch** wraps batched JAX arrays with dict-style access plus per-element state
- All modules inherit from `DataraxModule` for consistent behavior
- Protocols enable duck-typing with `isinstance()` checks

Architecture

DataraxModule (base)
├── OperatorModule          → Transformations (learnable)
└── StructuralModule        → Non-parametric structural processors
    ├── DataSourceModule    → Data loading
    ├── BatcherModule       → Batching logic
    └── SamplerModule       → Index sampling

Quick Start

import jax.numpy as jnp
from datarax.core import Element, Batch
from datarax.core.config import OperatorConfig
from datarax.core.metadata import Metadata

# Create an element
element = Element(
    data={"image": jnp.zeros((32, 32, 3))},
    state={"step": 0},
    metadata=Metadata(index=0),
)

# Access and update immutably
new_element = element.replace(
    data={"image": element.data["image"] / 255.0}
)

Modules

Data Structures

element_batch - Element and Batch data containers
metadata - Metadata handling and field selection
spec - Element specs: data as given or as JAX arrays, and batch validation
prng - The named nnx.Rngs streams and per-record key derivation

Configuration

config - Configuration base classes and validation

Base Classes

module - DataraxModule base class
operator - OperatorModule for transformations
data_source - DataSourceModule for data loading
batcher - BatcherModule for batch creation
sampler - SamplerModule for index sampling

Specialized

cross_modal - Cross-modal data handling
modality - Base classes for single-modality (per-field) operators
structural - Structural utilities and patterns

Real-World Examples

DADA Learned Augmentation - nnx.Param and nnx.value_and_grad through augmentation operators for policy search
Learned ISP Pipeline - Custom ModalityOperator subclasses with learnable parameters for ISP stages
DDSP Audio Synthesis - Custom OperatorModule subclasses for audio, showcasing extensibility to any domain