Utilities

Utility modules providing common functionality across Datarax. These are low-level helpers used internally and available for advanced use cases.

Available Utilities

Utility	Purpose	Key Functions
PyTree Utils	Batch tree operations	Batch size, split, concatenate
External	Library adapters	Integration helpers
Cache	Dataset cache layout	Resolve/structure cache dirs
Multirate	Signal alignment	Align streams at different rates

!!! note "Key points"

- PyTree utils operate on `Element`/`Batch` structures
- Use these for custom operators and extensions
- Randomness helpers live in the core layer: [prng](../core/prng.md)

Quick Start

from datarax.utils.pytree_utils import get_batch_size

# Inspect a batch produced by a pipeline
size = get_batch_size(batch)

Modules

• pytree_utils - Batch PyTree manipulation and transformation
• external - External library adapters and integrations
• cache - Dataset cache-layout helpers
• multirate - Multirate signal-alignment helpers

Batch PyTree Operations

from datarax.utils.pytree_utils import (
    add_batch_dimension,
    split_batch_for_devices,
    concatenate_batch_sequence,
)

# Promote a single element to a batch of size 1
batch = add_batch_dimension(element)

# Split a batch across devices, then recombine
shards = split_batch_for_devices(batch, num_splits=4)
merged = concatenate_batch_sequence(shards)

See Also

• Types & Protocols - Type definitions
• JAX Documentation - JAX PyTrees