AudioTree

AudioTree is an audio data loading and augmentation library supporting PyTorch and with extra features for JAX. Its central type, AudioTree, holds a batch of audio as a pytree: waveform, sample rate, loudness, and more, including your own per-item arrays. Augmentations come in matched NumPy and JAX backends (NumPy for CPU data-loader workers, JAX for jitted training steps).

The documentation carries the full guides, tested examples, and the API reference; start there.

pip install audiotree

JAX, Flax, Grain, NumPy, librosa, and soundfile come with it. Two extras: audiotree[bagz] adds Bagz record files, which back string leaves in TreeWriter/TreeDataSource and the windowed-LUFS cache (an extra rather than a dependency because bagz publishes manylinux x86-64 wheels only, and a hard dependency made pip install audiotree unsatisfiable elsewhere); audiotree[progress] adds tqdm for AudioWriter(show_progress=True). audiotree[all] is both, wherever bagz has a wheel.

from audiotree import AudioTree
from audiotree.sources import create_audio_dataset
from audiotree.transforms import stereo, volume_norm

# One file in, one AudioTree out. Even a single file is a batch (of 1).
audio = AudioTree.from_file("/data/audio/song.wav", sample_rate=44_100)
print(audio.waveform.shape)   # (1, channels, samples)

# The same idea for a whole directory: a shuffled, infinite stream of 5-second
# excerpts. num_epochs=None repeats forever; an integer gives that many passes
# over the files.
ds = create_audio_dataset(
    sources=["/data/audio"],
    sample_rate=44_100,
    duration=5.0,
    shuffle=True,
    num_epochs=None,
)

# Augment. One .seed() call: each random_map derives its own stream from it.
ds = ds.seed(42)
ds = ds.map(stereo())
ds = ds.random_map(volume_norm(min_db=-20, max_db=-15))

# Batch. AudioTree.batch concatenates along the leading axis items already have,
# rather than stacking a new one.
it = iter(ds.to_iter_dataset().batch(8, batch_fn=AudioTree.batch))

batch: AudioTree = next(it)
print(batch.waveform.shape)   # (8, 2, 220500) == (batch, channels, samples)
print(batch.sample_rate)      # 44100, one scalar for the whole batch
print(batch.lufs.shape)       # (8,): volume_norm leaves the achieved loudness behind
print(batch.filepath[0])      # the source file item 0 was drawn from

The same transforms exist in two backends: audiotree.transforms (NumPy, for CPU Grain workers) and audiotree.transforms.jax (JAX, for jitted training steps), bindable from YAML or the command line with DBraun's ArgBind fork. Beyond the quickstart: balanced sampling across source groups, length-aware windowed sampling, loudness-gated excerpt search, per-dataset read-error policies, two on-disk dataset writers, and neural-codec protocols. The guides cover all of it, including how AudioTree compares to audiotools and torchaudio.

AudioTree follows Effort-based Versioning: the version communicates the effort a change is likely to cost you, not a syntactic classification. Breaking changes are documented in the changelog.

@software{Braun_AudioTree_2026,
   author = {Braun, David},
   title = {{AudioTree}},
   url = {https://github.com/DBraun/audiotree},
   version = {1.0.0},
   year = {2026}
}

See CITATION.cff.

MIT, with third-party notices for the julius-derived resampler and the pyloudnorm-derived loudness code under LICENSES/. The audio fixtures under tests/assets/ in the repository are carved out of the MIT grant (the MUSDB18-HQ excerpt is CC BY-NC-SA 4.0) and are not part of any published distribution.