audiotree.transforms¶

AudioTree transforms for data augmentation.

This module provides two sets of transforms:

NumPy transforms (this module, audiotree.transforms): For CPU-based grain data pipelines. Uses NumPy operations and np.random.Generator.
JAX transforms (audiotree.transforms.jax): For GPU/JIT training pipelines. Uses JAX operations and jax.random.key.

Example - CPU (grain pipeline):

from audiotree.transforms import volume_norm, trim

# Create transform and apply with grain
transform = volume_norm(min_db=-20, max_db=-15)
ds = ds.random_map(transform, seed=42)

Example - GPU (jitted training step):

from audiotree.transforms import jax as jax_transforms
import argbind

# Bind all transforms for YAML configuration
transforms_lib = argbind.bind_module(jax_transforms)

@argbind.bind("train", "val")
def augment_batch(rng, batch, transforms: list[str] = None):
    for transform_name in transforms or []:
        transform = getattr(transforms_lib, transform_name)()
        if hasattr(transform, "random_map"):
            rng, subkey = jax.random.split(rng)
            batch = transform.random_map(batch, subkey)
        elif hasattr(transform, "map"):
            batch = transform.map(batch)
    return batch

See also

audiotree.transforms.jax for GPU transforms

class AudioCodec(*args, **kwargs)¶

Structural protocol for codecs that encode audio to discrete codes.

encode() takes an AudioTree and returns codes: an integer array whose shape convention (e.g. (batch, codebooks, frames)) is defined by the codec.

The codec is responsible for resampling the input to its own sample rate and for its channel handling (e.g. folding stereo into the batch for a mono encoder). Any state it needs beyond the codes – for example a loudness-normalization factor applied before quantizing – is the codec’s own to keep; audiotree stores only what encode returns.

This protocol is runtime_checkable, so isinstance(codec, AudioCodec) reports whether an object supplies encode. It is deliberately separate from LatentAudioCodec: a codec that only produces codes satisfies this protocol on its own, and one that does both satisfies both.

class LatentAudioCodec(*args, **kwargs)¶

Structural protocol for codecs that encode audio to continuous latents.

encode_to_latent() takes an AudioTree and returns a latent array, whose shape convention is defined by the codec. As with AudioCodec, the codec owns resampling and channel handling, and the protocol is runtime_checkable.

class choose(*transforms, c: int = 1, weights=None, prob: float = 1.0)¶

Choose c transform(s) among transforms with optional probability weights.

With probability prob, choose c transform(s) from the list of transforms and apply them sequentially.

NumPy backend only. Which transforms run is decided in Python, so this cannot be traced by jax.jit; there is deliberately no audiotree.transforms.jax.choose. Compose JAX transforms explicitly instead.

This is a hand-written grain.transforms.RandomMap, not a decorated transform, so it does not take split_seed, scope or output_key; scope the transforms handed to it instead. Its prob is also a single draw for the whole element, not one per batch item as it is for a decorated transform.

Parameters:

*transforms – Transforms to choose from. Each must be a grain.transforms.Map or grain.transforms.RandomMap — which is what every audiotree transform constructor returns.
c – Number of transforms to choose
weights – Optional probability weights for each transform. Must be one weight per transform, each non-negative, summing to 1.
prob – Probability of applying any transforms at all

Raises:

TypeError – If a positional argument is not a grain transform.
ValueError – If c, weights or prob are out of range.

Example:

transform = choose(
    volume_change(min_db=-6, max_db=6),
    invert_phase(),
    swap_stereo(),
    c=2,
    weights=[0.5, 0.3, 0.2],
    prob=0.9,
)
ds = ds.random_map(transform, seed=42)

random_map(element, rng: Generator)¶: Maps a single element.

corrupt_phase(amount: float = 1.0, hop_factor: float = 0.5, frame_length: int = 2048, window: str = 'hann', keep_lufs: bool = False, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Perform phase corruption on audio.

The phase shift range is [-pi * amount, pi * amount], independently selected for each channel and frequency of the STFT, and shared across frames. Contrast shift_phase(), which rotates the whole spectrum of an item by one angle.

Parameters:

amount – Maximum phase shift in multiples of pi (0.0 to 1.0)
hop_factor – Hop size as fraction of frame_length, in (0, 0.5]. Larger hops leave the analysis windows unable to reconstruct the signal, so they are rejected.
frame_length – STFT frame length in samples
window – Window function name
keep_lufs – If True, preserve the cached lufs and lufs_windows. Phase corruption leaves the magnitude spectrum (and thus energy) intact, so loudness is approximately unchanged; the cached values are invalidated by default to be safe.
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with corrupted phase

Raises:

ValueError – If amount is negative, hop_factor is outside (0, 0.5], or the audio is shorter than frame_length samples.

Example

transform = corrupt_phase(amount=0.5, hop_factor=0.5) ds = ds.random_map(transform, seed=42)

Create a transform that encodes audio to continuous latents.

Calls codec.encode_to_latent(audio_tree) and stores the result on AudioTree.latents. AudioTrees that already have latents pass through unchanged.

Parameters:

codec – Object implementing encode_to_latent(AudioTree) -> latents (see LatentAudioCodec).
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them).
output_key – Write the result under a new key instead of replacing the input.

Returns:

Transform for use with .map().

Example

transform = encode_latents(codec) ds = ds.map(transform)

Create a transform that encodes audio to discrete codes.

Calls codec.encode(audio_tree) and stores the returned codes on AudioTree.codes exactly as the codec produced them (the codec defines the shape convention). AudioTrees that already have codes pass through unchanged.

Parameters:

codec – Object implementing encode(AudioTree) -> codes (see AudioCodec).
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them).
output_key – Write the result under a new key instead of replacing the input.

Returns:

Transform for use with .map().

Example

transform = encode_with_codec(codec) ds = ds.map(transform)

identity(*, scope=None, output_key=None)¶

Return audio without any modifications.

Useful as a placeholder or for testing.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree unchanged

Example

transform = identity() ds = ds.map(transform)

invert_phase(*, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Invert the phase of all channels of audio.

For data augmentation, it’s common to use prob=0.5 to apply this transform probabilistically.

Parameters:

prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with inverted phase

Example

transform = invert_phase(prob=0.5) ds = ds.random_map(transform, seed=42)

map_transform(fn: Callable) → Callable¶

Decorator to create a MapTransform from a function.

The decorated function should have signature:

fn(audio_tree: AudioTree, **params) -> AudioTree

The returned callable constructs the transform. It takes the function’s own parameters (positionally or by keyword) plus the keyword-only scope and output_key. A misspelled parameter raises TypeError rather than being silently ignored, as does omitting one that has no default. Parameter values are opaque, so a dict, a list or None is passed through to the function unchanged.

Usage:

@map_transform
def trim(audio_tree, length=1.0):
    ...
    return audio_tree

transform = trim(length=3.0)
ds = ds.map(transform)

mono(*, scope=None, output_key=None)¶

Convert audio to mono by averaging channels.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with mono audio

Example

transform = mono() ds = ds.map(transform)

peak_norm(*, scope=None, output_key=None)¶

Peak-normalize audio so the largest absolute value is 1.0.

Unlike rescale_audio(), which only scales down audio that exceeds the [-1.0, 1.0] range, this always divides by the peak so the result peaks at 1.0. The peak is computed per item in the batch (across channels and samples) and clamped to a small epsilon to avoid division by zero on silent audio.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with peak-normalized audio

Example

transform = peak_norm() ds = ds.map(transform)

random_transform(fn: Callable) → Callable¶

Decorator to create a RandomMapTransform from a function.

The decorated function should have signature:

fn(audio_tree: AudioTree, rng, **params) -> AudioTree

The returned callable constructs the transform. It takes the function’s own parameters (positionally or by keyword) plus the keyword-only prob, split_seed, scope and output_key. A misspelled parameter raises TypeError rather than being silently ignored, as does omitting one that has no default. Parameter values are opaque, so a dict, a list or None is passed through to the function unchanged.

Usage:

@random_transform
def volume_norm(audio_tree, rng, min_db=-20.0, max_db=-15.0):
    ...
    return audio_tree

transform = volume_norm(min_db=-30, max_db=-10, prob=0.9)
ds = ds.random_map(transform, seed=42)

resample(sample_rate: int | None = None, *, scope=None, output_key=None)¶

Resample audio to a new sample rate.

Wraps resample(): NumPy-backed waveforms resample on CPU via librosa, JAX-backed waveforms use the JAX/Julius port.

Parameters:

sample_rate – Target sample rate in Hz (e.g. 16000). Required.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree resampled to sample_rate

Example

transform = resample(sample_rate=16000) ds = ds.map(transform)

rescale_audio(*, scope=None, output_key=None)¶

Rescale audio so the largest absolute value is 1.0.

If all values are already in [-1.0, 1.0], no transformation is applied. Useful if transforms have caused the audio to clip.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with rescaled audio

Example

transform = rescale_audio() ds = ds.map(transform)

roll(min_seconds: float = 0.0, max_seconds: float = 0.0, mode: str = 'wrap', *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Apply a circular shift (roll) to audio data.

The amount of roll is randomly selected per item in the batch (not per channel). Positive values roll the audio to the right, negative values roll to the left.

Rolling invalidates the recorded source-file offset provenance (it no longer says where sample 0 came from). With prob < 1 the offset is dropped for the whole batch, not just the rolled items: a per-item mix of “valid” and “invalidated” cannot be represented in one array, and keeping stale offsets on the rolled items would be worse.

Parameters:

min_seconds – Minimum roll amount in seconds (negative = left)
max_seconds – Maximum roll amount in seconds (positive = right)
mode – Padding mode - “wrap” (circular) or “constant” (zero padding)
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with rolled audio

Example

transform = roll(min_seconds=-1.0, max_seconds=1.0, mode=”wrap”) ds = ds.random_map(transform, seed=42)

shift_phase(amount: float = 1.0, keep_lufs: bool = False, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Perform a phase shift on audio.

The phase shift range is [-pi * amount, pi * amount]. One angle is drawn per item in the batch and applied to every channel and frequency, so a stereo image stays coherent. Contrast corrupt_phase(), which draws an angle per channel and frequency.

Parameters:

amount – Maximum phase shift in multiples of pi
keep_lufs – If True, preserve the cached lufs and lufs_windows. A phase shift leaves the magnitude spectrum (and thus energy) intact, so loudness is approximately unchanged; the cached values are invalidated by default to be safe.
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with shifted phase

Raises:

ValueError – If amount is negative, or the audio is shorter than one STFT frame (2048 samples).

Example

transform = shift_phase(amount=0.5) ds = ds.random_map(transform, seed=42)

stereo(*, scope=None, output_key=None)¶

Convert audio to stereo by duplicating mono channel.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with stereo audio

Example

transform = stereo() ds = ds.map(transform)

swap_stereo(*, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Exchange the left and right channels of stereo audio.

Mono audio passes through unchanged: with one channel the only possible permutation is the identity. Audio with three or more channels raises, because which pair to exchange is undefined.

For data augmentation, it’s common to use prob=0.5 to apply this transform probabilistically.

Parameters:

prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with swapped channels

Raises:

ValueError – If the audio has more than two channels.

Example

transform = swap_stereo(prob=0.5) ds = ds.random_map(transform, seed=42)

trim(length: float = 1.0, mode: str = 'wrap', *, scope=None, output_key=None)¶

Adjust audio length to a fixed length in seconds.

If audio is shorter than the target length, it will be padded according to mode. If audio is longer, it will be trimmed.

Parameters:

length – Target length in seconds
mode – Padding mode if audio needs to be lengthened - “wrap”: Circular shift (default). Audio wraps around. - “constant”: Zero padding.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree adjusted to target length

Example

# Trim to 3 seconds transform = trim(length=3.0) ds = ds.map(transform)

# Pad short audio with zeros transform = trim(length=5.0, mode=”constant”) ds = ds.map(transform)

volume_change(min_db: float = 0.0, max_db: float = 0.0, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Change the volume by a uniformly randomly selected decibel value.

Parameters:

min_db – Minimum gain change in dB
max_db – Maximum gain change in dB
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with volume changed

Raises:

ValueError – If min_db > max_db.

Example

transform = volume_change(min_db=-12, max_db=12, prob=0.9) ds = ds.random_map(transform, seed=42)

volume_norm(min_db: float = 0.0, max_db: float = 0.0, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Normalize volume to a randomly selected loudness value specified in LUFS.

A tree arriving with lufs already populated (e.g. restored from a written manifest) is trusted and not re-measured; loudness is measured only when lufs is unset. The cache is trustworthy because every operation that changes the audio clears it.

Parameters:

min_db – Minimum target loudness in LUFS
max_db – Maximum target loudness in LUFS
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with normalized loudness

Raises:

ValueError – If min_db > max_db.

Example

transform = volume_norm(min_db=-20, max_db=-15) ds = ds.random_map(transform, seed=42)

audiotree.transforms.jax¶

JAX-based transforms for GPU/JIT training pipelines.

Use these transforms inside @jax.jit functions or when working with JAX arrays on GPU. These transforms use JAX operations and accept jax.random.key for random transforms.

For CPU-based grain data pipelines, use audiotree.transforms instead.

Example - Direct usage:

from audiotree.transforms import jax as jax_transforms
import jax

# Create transform and apply with JAX key
transform = jax_transforms.volume_norm(min_db=-20, max_db=-15)
rng = jax.random.key(42)
result = transform.random_map(audio, rng)

Example - ArgBind-configured training pipeline:

from audiotree.transforms import jax as jax_transforms
import argbind

# Bind all transforms for YAML configuration
transforms_lib = argbind.bind_module(jax_transforms)

@argbind.bind("train", "val")
def augment_batch(rng, batch, transforms: list[str] = None):
    for transform_name in transforms or []:
        transform = getattr(transforms_lib, transform_name)()
        if hasattr(transform, "random_map"):
            rng, subkey = jax.random.split(rng)
            batch = transform.random_map(batch, subkey)
        elif hasattr(transform, "map"):
            batch = transform.map(batch)
    return batch

See also

audiotree.transforms for CPU/grain transforms

Perform phase corruption on audio.

Parameters:

amount – Maximum phase shift in multiples of pi (0.0 to 1.0)
hop_factor – Hop size as fraction of frame_length, in (0, 0.5]. Larger hops leave the analysis windows unable to reconstruct the signal, so they are rejected.
frame_length – STFT frame length in samples
window – Window function name
keep_lufs – If True, preserve the cached lufs and lufs_windows. Phase corruption leaves the magnitude spectrum (and thus energy) intact, so loudness is approximately unchanged; the cached values are invalidated by default to be safe.
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with corrupted phase

Raises:

ValueError – If amount is negative, hop_factor is outside (0, 0.5], or the audio is shorter than frame_length samples.

Create a transform that encodes audio to continuous latents.

Calls codec.encode_to_latent(audio_tree) and stores the result on AudioTree.latents. AudioTrees that already have latents pass through unchanged.

Parameters:

codec – Object implementing encode_to_latent(AudioTree) -> latents (see LatentAudioCodec).
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them).
output_key – Write the result under a new key instead of replacing the input.

Returns:

Transform for use with .map().

Example

transform = encode_latents(codec) ds = ds.map(transform)

Create a transform that encodes audio to discrete codes.

Parameters:

codec – Object implementing encode(AudioTree) -> codes (see AudioCodec).
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them).
output_key – Write the result under a new key instead of replacing the input.

Returns:

Transform for use with .map().

Example

transform = encode_with_codec(codec) ds = ds.map(transform)

identity(*, scope=None, output_key=None)¶

Return audio without any modifications.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

invert_phase(*, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Invert the phase of all channels of audio.

Parameters:

prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with inverted phase

map_transform(fn: Callable) → Callable¶

Decorator to create a MapTransform from a function.

The decorated function should have signature:

fn(audio_tree: AudioTree, **params) -> AudioTree

Usage:

@map_transform
def trim(audio_tree, length=1.0):
    ...
    return audio_tree

transform = trim(length=3.0)
ds = ds.map(transform)

mono(*, scope=None, output_key=None)¶

Convert audio to mono by averaging channels.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

peak_norm(*, scope=None, output_key=None)¶

Peak-normalize audio so the largest absolute value is 1.0.

Unlike rescale_audio(), which only scales down audio that exceeds the [-1.0, 1.0] range, this always divides by the peak (clamped to a small epsilon) so the result peaks at 1.0. The peak is computed per item in the batch, across channels and samples.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

random_transform(fn: Callable) → Callable¶

Decorator to create a RandomMapTransform from a function.

The decorated function should have signature:

fn(audio_tree: AudioTree, rng, **params) -> AudioTree

Usage:

@random_transform
def volume_norm(audio_tree, rng, min_db=-20.0, max_db=-15.0):
    ...
    return audio_tree

transform = volume_norm(min_db=-30, max_db=-10, prob=0.9)
ds = ds.random_map(transform, seed=42)

resample(sample_rate: int | None = None, *, scope=None, output_key=None)¶

Resample audio to a new sample rate (JAX/Julius port for JAX waveforms).

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

rescale_audio(*, scope=None, output_key=None)¶

Rescale audio so the largest absolute value is 1.0.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

roll(min_seconds: float = 0.0, max_seconds: float = 0.0, mode: str = 'wrap', *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Apply a circular shift (roll) to audio data.

Parameters:

min_seconds – Minimum roll amount in seconds (negative = left)
max_seconds – Maximum roll amount in seconds (positive = right)
mode – Padding mode - “wrap” (circular) or “constant” (zero padding)
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with rolled audio

shift_phase(amount: float = 1.0, keep_lufs: bool = False, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Perform a phase shift on audio.

Parameters:

amount – Maximum phase shift in multiples of pi
keep_lufs – If True, preserve the cached lufs and lufs_windows. A phase shift leaves the magnitude spectrum (and thus energy) intact, so loudness is approximately unchanged; the cached values are invalidated by default to be safe.
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with shifted phase

Raises:

ValueError – If amount is negative, or the audio is shorter than one STFT frame (2048 samples).

stereo(*, scope=None, output_key=None)¶

Convert audio to stereo by duplicating mono channel.

Parameters:

scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

swap_stereo(*, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Exchange the left and right channels of stereo audio.

Mono audio passes through unchanged: with one channel the only possible permutation is the identity. Audio with three or more channels raises, because which pair to exchange is undefined.

Parameters:

prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with swapped channels

Raises:

ValueError – If the audio has more than two channels.

trim(length: float = 1.0, mode: str = 'wrap', *, scope=None, output_key=None)¶

Adjust audio length to a fixed length in seconds.

Parameters:

length – Target length in seconds
mode – Padding mode - “wrap” (circular) or “constant” (zero padding)
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree adjusted to target length

volume_change(min_db: float = 0.0, max_db: float = 0.0, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Change the volume by a uniformly randomly selected decibel value.

Parameters:

min_db – Minimum gain change in dB
max_db – Maximum gain change in dB
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with volume changed

Raises:

ValueError – If min_db > max_db.

volume_norm(min_db: float = 0.0, max_db: float = 0.0, *, prob=1.0, split_seed=True, scope=None, output_key=None)¶

Normalize volume to a randomly selected loudness value specified in LUFS.

Parameters:

min_db – Minimum target loudness in LUFS
max_db – Maximum target loudness in LUFS
prob – Probability of applying the transform, drawn independently per batch item. Defaults to 1.0 (always).
split_seed – Give each AudioTree leaf its own RNG split. With False every leaf draws identically, which keeps a dry/wet pair in lockstep. Defaults to True.
scope – Which leaves of a dict-of-AudioTree element to transform. Defaults to None (all of them); an explicitly empty scope ([] or {}) selects none, making the transform a no-op. See Working with Dict[str, AudioTree] Batches.
output_key – Write the result under a new key instead of replacing the input. A plain string names exactly one output, so it requires a single in-scope leaf; with several, pass a callable that maps each leaf’s path to a distinct name.

Returns:

AudioTree with normalized loudness

Raises:

ValueError – If min_db > max_db.