Core Components

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

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)
├── SharderModule           → Device sharding
└── 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

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
• sharder - SharderModule for device sharding

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

See Also

• Types & Protocols - Type definitions
• Configuration Guide - Detailed config documentation
• DAG Executor - Using core components in pipelines