Skip to content

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

★ Insight ─────────────────────────────────────

  • Element wraps a single data sample with state and metadata
  • Batch is a dictionary of batched JAX arrays
  • All modules inherit from DataraxModule for consistent behavior
  • Protocols enable duck-typing with isinstance() checks

─────────────────────────────────────────────────

Architecture¤

DataraxModule (base)
├── OperatorModule      → Transformations (learnable)
├── DataSourceModule    → Data loading
├── BatcherModule       → Batching logic
├── SamplerModule       → Index sampling
└── SharderModule       → Device sharding

Quick Start¤

from datarax.core import Element, Batch
from datarax.core.config import OperatorConfig

# 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¤

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 - Data modality definitions (image, text, audio)
  • structural - Structural utilities and patterns

Real-World Examples¤

See Also¤