Changelog

The format is based on Keep a Changelog (categories: Added, Changed, Deprecated, Removed, Fixed, Security), and AudioTree follows Effort-based Versioning.

[Unreleased]

Nothing yet.

[1.0.0] - 2026-08-11

The first stable release, and a near-total rewrite of the public surface. 0.2.x was a handful of grain data sources with class-based transforms; 1.0 is a typed AudioTree container with two transform backends (NumPy for grain workers, JAX for jitted training loops), four dataset constructors, and two writer/reader pairs with versioned on-disk formats. The public surface is the set of names exported from the audiotree, audiotree.sources, audiotree.transforms and audiotree.transforms.jax namespaces, pinned by tests/test_public_api.py.

Assume nothing from 0.2.x still compiles. The field audio_data is now waveform, loudness is now lufs (with replace_lufs() / normalize_lufs() / lufs_cutoff / keep_lufs= to match), and the merged metadata dict is split — payload keys are extras, while the filepath/source provenance moves to a private _metadata container (spelled metadata on disk); every class-based transform is a snake_case function (VolumeNorm to volume_norm); AudioDataSimpleSource / AudioDataBalancedSource / AudioDataBalancedDataset are create_audio_dataset() and create_balanced_audio_dataset(); SaliencyParams is ExcerptConfig; the NumPy transforms take an np.random.Generator rather than a jax.random.key; and public callables are keyword-only past their first argument or two.

Two changes are silently corpus-affecting – they alter the data a pipeline produces rather than raising, so re-render anything pre-rendered and re-check anything trained. The dataset loaders used to read the first duration seconds of every file on every epoch (the default is now a random offset), and the loudness meter under-reported excerpts shorter than 400 ms and used two different K-weighting filters within one call. Each is detailed below.

Requires Python 3.11-3.14 on Linux or macOS; Windows works except for grain’s multiprocessing prefetch (google/grain#793). Artifacts written by a pre-1.0 development build carry no format header and are refused with a “re-render the dataset” message – none of those layouts ever shipped in a release, so there is no released data to migrate.

Breaking changes at a glance. Renames: waveform, lufs, extras, filepath=, predicate=, shuffle_seed= / excerpt_seed=, ExcerptConfig, audiotree.sources. Signatures: keyword-only past the leading positionals; transform constructors reject unknown names; scope=["wet"] replaces scope={"wet": True}. Behaviour: the loudness meter, split_seed=False, prob<1 and the per-group seeds all changed what they produce. Removals: AudioTree.from_array, ReduceBatchTransform, num_records on the dataset builders, and the post-mix shuffle.

Added

• ExcerptConfig: The one object that says how an excerpt is chosen, exported from both audiotree and audiotree.sources and accepted as excerpt= by create_audio_dataset(), create_balanced_audio_dataset(), AudioTree.excerpt() and AudioTree.loudest_excerpt(). A frozen dataclass with one named strategy — "random" (a uniform offset; the default), "start" (offset 0) or "loudest" (best-of-num_tries search for a window clearing lufs_cutoff, default −40 LUFS) — plus search ("uniform", "bias_early", or a dotted path to your own; it draws the offsets of both "random" and "loudest", and only "start" — which draws nothing — rejects it) and on_failure ("keep" / "skip" / "raise"). Setting a knob the chosen strategy ignores raises at construction rather than being silently dropped. It replaces SaliencyParams; see Changed for the migration. The two built-in searchers live at audiotree.core.search_uniform / search_bias_early, which are internal — refer to them by name (search="bias_early") rather than importing them.

• AudioTree.backend and AudioTree.device: Which array library a tree’s leaves belong to ("numpy", "jax", or "mixed") and which JAX device they sit on (None on NumPy). A tree stops being homogeneous more easily than you would expect — a NumPy-namespace transform applied to a JAX tree hands back a NumPy waveform — and the failure is silent until jax.jit starts re-uploading a leaf every call. backend reports "mixed" rather than raising, so it is safe to log; device raises when the leaves genuinely disagree, because there is no honest single answer. jax.device_put / jax.device_get move a tree; these say where it currently is.

• TreeDataSource releases its memmaps: it is a context manager with an explicit close(), so with TreeDataSource(d) as src: ... (or a manual close()) drops every open mapping. It is a release, not a teardown — reading again reopens transparently, close() is idempotent, and a read in flight on another thread (grain’s prefetch pool) keeps its own snapshot of the handles, so it completes whole rather than silently returning an empty sample. This matters on Windows, where an open memmap blocks deletion: a run that read a dataset and then tried to remove it failed with PermissionError: [WinError 32] and nothing naming the holder.

• Optional extras: audiotree[bagz] for string leaves in TreeWriter/TreeDataSource and the windowed-LUFS cache; audiotree[progress] for AudioWriter(show_progress=True); audiotree[all] for both. [all] guards bagz behind environment markers so it stays installable everywhere, while an explicit pip install audiotree[bagz] fails loudly on the platforms bagz publishes no wheel for (see Changed).

• create_windowed_audio_dataset(): Length-aware window sampling that changes the unit of sampling from the file to a fixed-duration window. Each file is tiled into m_i = round(n_i ** alpha) evenly-spaced slots (where n_i is its number of hop windows), all slots are flattened into one globally-shuffled index, and each draw loads a jittered excerpt from its slot. This gives even per-file coverage, length-proportional sampling, and diverse batches (no lingering on long files) simultaneously. The alpha knob (in [0, 1]) tunes length bias: 0 = uniform per file, 1 = proportional to length, 0.5 = sqrt(length). Jittered offsets are confined to each file so windows never run past end-of-file. Exported from audiotree.sources.

• WindowConfig for balanced datasets: create_balanced_audio_dataset() accepts a window=WindowConfig(...) that builds every file-based group with create_windowed_audio_dataset() instead of one-excerpt-per-file. Group weights balance across categories while alpha controls length bias within each, composing multiplicatively. Mutually exclusive with a customized excerpt.

• Windowed-loudness cache (bagz): build_window_lufs_cache() runs a one-time preprocessing pass that computes per-file windowed integrated loudness (ITU-R BS.1770) and persists the ragged per-file LUFS arrays as Bagz records (no padding), plus durations and the measurement sample_rate/mono, via save_window_lufs() / load_window_lufs(). Passing lufs_cache= to create_windowed_audio_dataset() (or via WindowConfig) drops windows below lufs_cutoff at build time, turning saliency into a one-time filter rather than a per-visit search. Loudness is computed on the CPU with the upstream loudness library (loudness.integrated_loudness, the same kernel AudioTree.replace_lufs() uses for the integrated lufs of a NumPy waveform) — no JAX/GPU work, so the preprocessing is safe to run before forking grain workers and keeps the data-source layer free of GPU computation. The dataset verifies the cache’s sample_rate/mono match the audio it loads and refuses to run on a mismatch; the cache also records each file’s byte size, so a file re-rendered in place after the build names itself at dataset construction instead of silently keeping the stale durations and loudness, and a file the cache has never seen (added after the build) is reported with the rebuild remedy rather than a bare KeyError. Curated corpora can skip this entirely (the default).

• scan_durations(): Header-only (no decode) per-file duration scan, cacheable and passable as durations= so alpha/hop/duration retune without re-reading the corpus.

• peak_norm() transform: Peak-normalizes audio so its largest absolute value is 1.0, dividing by the per-item peak (across channels and samples, clamped to a small epsilon). Unlike rescale_audio(), which only scales down audio exceeding [-1.0, 1.0], this always normalizes to a peak of 1.0. Available in both the NumPy (audiotree.transforms) and JAX (audiotree.transforms.jax) backends.

• TreeWriter and TreeDataSource: A pytree-native writer and reader for memory-mapped datasets, both new in 1.0. Instead of manually extracting and reconstructing AudioTree fields, TreeWriter uses jax.tree_util to decompose any pytree into leaves — each array leaf becomes its own memmap .bin, each string leaf its own Bagz file — and a JSON structure descriptor in manifest.json enables exact reconstruction. It accepts AudioTree objects, dicts of arrays, dicts of AudioTrees, or nested combinations, all through a single write() method (call open() first, or use the writer as a context manager). The reader takes a directory path and implements Grain’s RandomAccessDataSource with a minimal interface (__len__, __getitem__).

• String leaf support in TreeWriter / TreeDataSource: write() now accepts str or List[str] leaves alongside arrays and AudioTrees. Strings are stored in Bagz files (one per string leaf) with no length limits or truncation. Single-sample reads return a bare str, and AudioTree.batch naturally collects them into List[str]. Backed by Bagz, which is an optional extra: install audiotree[bagz] to use string leaves. Datasets without string leaves need nothing extra, and exclude_prefixes can skip the string leaves of a dataset written elsewhere.

• Selective field loading via exclude_prefixes: TreeDataSource accepts an exclude_prefixes parameter to skip loading specific leaves by dot-separated name prefix. For example, exclude_prefixes=["wet.waveform"] skips the audio memmap, and exclude_prefixes=["dry"] skips all leaves under the dry subtree. Excluded leaves are omitted from the reconstructed pytree (AudioTree fields default to None, extras keys are absent). Prefix matching uses exact-or-dot semantics ("dry" matches "dry" and "dry.waveform" but not "dryness"). Excluding the root leaf of a bare-array dataset (a dataset whose whole pytree is one array) leaves nothing to return and raises at construction, rather than handing the caller an internal sentinel object.

• load_into_memory for worker-shared data: TreeDataSource accepts load_into_memory=True to load all non-excluded array and string leaves into RAM at init time, so the source never opens a file again — not in this process and not in a worker. The store survives pickling, so it reaches grain’s mp_prefetch workers (which are spawned, not forked) by pickle, and a forked child by copy-on-write; either way the workers do no disk I/O. String leaves (bagz) are read into a List[str]. __getitem__ copies each sample out of the store rather than handing out a view, so one caller’s in-place write cannot rewrite the dataset for every reader after it.

• Graceful undershoot handling in TreeWriter: expected_samples is an allocation hint rather than an exact requirement. On close(), if fewer samples were written than allocated, the memmap files are truncated to the actual sample count via os.truncate(), so an over-estimate costs no disk space and the reader still sees exactly the samples that were written. This avoids errors when the exact dataset size isn’t known ahead of time (e.g., split-dependent counts). Overshoot is governed separately by the keyword-only on_overflow=: "grow" (default) reallocates every leaf and keeps writing, "error" raises naming the written, allocated and offered counts, and "trim" writes what fits and warns. Like AudioWriter, the writer refuses a directory that already holds a dataset unless exist_ok=True.

• AudioTree.replace_extras(): Merges keyword arguments into extras — syntactic sugar for tree.replace(extras={**tree.extras, **kwargs}). Passed keys overwrite same-named existing keys, all other keys are kept, and neither the original tree nor its extras dict is mutated. AudioTree.replace() itself stays strict (unknown field names still raise), so typos don’t silently become extras entries.

• AudioTree.write(): Saves a single-item tree (batch_size == 1) to an audio file via soundfile, the inverse of from_file(). Takes a filepath (format inferred from its extension) plus optional keyword-only subtype/format/endian passthroughs and returns the Path written; it does not take a sample rate (uses self.sample_rate — call resample() first to change it). It raises ValueError — not a strippable assert — when the tree has no waveform, when it is mini-batched (rank 4), or when batch_size != 1, so index or iterate the batch first (e.g. for item in tree: item.write(...)).

• AudioTree.__iter__: AudioTree is now a proper collections.abc.Iterable. Iterating yields one batch-of-1 AudioTree per item along the leading batch axis (e.g. a batch-16 tree yields 16 trees of batch_size == 1). for loops worked before via the sequence protocol, but isinstance(tree, Iterable) now returns True.

• AudioTree.samples property: Returns the number of samples in the waveform (waveform.shape[-1]).

• AudioTree.source property: Returns the source group name for each item in the batch. When using create_balanced_audio_dataset() with sources={"music": [...], "speech": [...]}, each loaded AudioTree tracks which source group it came from.

• AudioTree.lufs_windows + per-window loudness: replace_lufs() now also fills a new lufs_windows field alongside lufs — the ungated K-weighted loudness of each analysis window (a loudness-over-time curve), shaped (*batch, num_windows). Unlike lufs (the gated integrated program loudness), every window is directly comparable and a fully silent window reads -inf, so a statistic like “fraction of windows above −30 LUFS” is meaningful. replace_lufs(lufs_window_sec=0.4, lufs_hop_sec=None) controls the window length (≥ 0.4s, the EBU momentary integration time) and the step between windows (defaults to lufs_window_sec for non-overlapping windows; a smaller hop overlaps them) — the same spelling WindowConfig and create_windowed_audio_dataset() use for the same quantities. The trailing partial window is dropped, so an excerpt shorter than one window yields an empty lufs_windows. The windowed curve is computed natively: NumPy uses exact ITU-R BS.1770 K-weighting IIR biquads via scipy on the CPU, and JAX uses jaxloudnorm’s FIR-approximated K-weighting so the whole batch stays on the accelerator (the integrated lufs still comes from the loudness library on the NumPy path and a vmapped jaxloudnorm meter on the JAX one). Where the measurement runs and which kernel runs are separate keyword-only arguments: device= (None / "cpu" / "gpu" / "tpu", mirroring jax.jit’s backend, or a jax.Device) and engine= (None / "numpy" / "jax"). replace_lufs(device="gpu") runs the vmapped jaxloudnorm kernel on the accelerator even for a NumPy waveform (much faster for large batches, since the CPU meter measures one item at a time) and still returns NumPy loudness, so you no longer need a manual jax.device_put/jax.device_get round-trip. None for both (the default) uses the waveform’s own array library on its current device. lufs_windows is a first-class field: it survives create() / from_file(), indexing, batching, and round-trips through AudioWriter/from_manifest and TreeWriter/TreeDataSource. Transforms keep it consistent — volume_norm/volume_change shift it by the applied gain, phase transforms preserve it under keep_lufs=True, and length/energy/channel changes invalidate it. For hand-rolled waveform edits, AudioTree.clear_lufs() is the documented shorthand for dropping both cached fields (replace(lufs=None, lufs_windows=None)) so the next replace_lufs() measures the new audio.

• AudioTree.normalize_lufs(target_lufs): Normalizes audio to a target LUFS level. Computes loudness if not already set, scales audio, and updates the lufs and lufs_windows fields (a constant gain shifts every window’s LUFS equally).

• AudioTree.create(): New classmethod with automatic audio dimensionality handling (1D→3D, 2D→3D) and filepath parameter support.

• AudioTree.filter(): Takes a predicate(AudioTree) -> bool applied to each batch item and returns a new AudioTree holding only the items it accepted.

• AudioTree.split(), AudioTree.reshape_mini_batches(), AudioTree.flatten_mini_batches(): Split an AudioTree into a list, and add or remove a mini-batch axis (useful for nnx.scan).

• source parameter on AudioTree.create(): create() now accepts a source= kwarg (encoded into the library-managed metadata provenance container and read back via the source property), matching from_file(). Unlike from_file(), which takes only a single string, create() accepts either a single string (tags the whole batch) or a list of strings with one source-group name per batch item — so a batched tree assembled outside the dataloaders (e.g. windows rendered from mixed materials) can still carry a correct per-item .source. Round-trips through TreeWriter/TreeDataSource.

• pad_mode parameter: Added to AudioTree.from_file(). Options include "constant" (zero padding), None (don’t pad), and "wrap" (loop audio).

• JAX transforms module: audiotree.transforms.jax provides JAX-native transforms for GPU/JIT training pipelines. Uses jax.random.key and JAX operations throughout. It exports volume_norm, volume_change, invert_phase, swap_stereo, corrupt_phase, shift_phase, roll, trim, mono, stereo, rescale_audio, peak_norm, resample, identity, the shared encode_with_codec / encode_latents codec transforms, and the @random_transform / @map_transform decorators. choose() is the one transform that stays NumPy-only.

• resample() transform: Resamples to a new sample rate in both the NumPy (audiotree.transforms) and JAX (audiotree.transforms.jax) backends, wrapping AudioTree.resample (so NumPy waveforms use librosa and JAX waveforms use the Julius port).

• Transform decorators: New @random_transform and @map_transform for building custom transforms from plain functions. They live in audiotree.transforms.decorators and are re-exported from audiotree.transforms and audiotree.transforms.jax, which is where you should import them from. @map_transform handles the boilerplate for the keyword-only scope and output_key parameters; @random_transform adds prob and split_seed on top and threads the RNG through.

• create_audio_dataset(): For creating simple audio datasets without balancing. Loads all files from one or more directories and applies grain’s random_map for proper RNG seeding. num_epochs: int | None sets the pass count over the corpus (None for an unbounded stream; defaults to 1 here and to None for the balanced builder).

• find_audio_files(): New public helper in audiotree.sources for listing audio files under one or more directories (skipping hidden files/dirs, filtering by extension). Each source may also be a glob pattern — any entry containing *, ?, or [...] is expanded, with each match then recursed into (if a directory) or kept (if a file), so patterns like "/mnt/d/musdb18hq/train/*/mixture.wav" and recursive "/data/**/*.wav" work directly. The result is sorted and de-duplicated, so the file order — and therefore seeded shuffling — is deterministic across machines and filesystems (previously raw os.walk order, which varied by filesystem).

• Saliency-aware excerpt loading: create_audio_dataset() and create_balanced_audio_dataset() choose each excerpt inside grain’s random_map, so a file visited again on a later epoch gets a fresh offset instead of the same one. How the offset is chosen is configured with excerpt=ExcerptConfig(...) ("random", "start", or a best-of-num_tries "loudest" search), and AudioTree.loudest_excerpt(path, rng, excerpt=...) exposes that search directly on the container.

• AudioWriter: New class for writing AudioTree batches to disk as individual audio files plus a manifest.npz index — one row per item carrying the filepath/source provenance columns, lufs, the other AudioTree label fields, and any extras_* and tag columns. The manifest is rewritten atomically (temp file plus os.replace) every manifest_every items (default 1000), so a render killed partway still leaves a readable prefix. The writer adopts the sample rate of the first AudioTree written and raises ValueError if a later write differs, so one manifest never mixes rates; the default encoding subtype is the widest the container supports (FLOAT for WAV/AIFF/CAF/W64/RF64, PCM_24 for FLAC) rather than libsndfile’s PCM_16, and an explicitly requested fixed-point subtype that clips raises a RuntimeWarning naming the worst offender; and it refuses to write into a directory that already holds a dataset unless exist_ok=True, so a finished render is never silently overwritten. write() commits per batch: every entry is validated before any audio lands, a mid-batch failure removes the files that call had written (no manifest row points at them) and leaves the index and manifest untouched, so a failed write() can be fixed and retried without duplicating items or orphaning audio.

• AudioDataSource: New class for reading AudioWriter outputs as a grain RandomAccessDataSource with metadata preservation and filtering capabilities. AudioDataSource(path) accepts the dataset directory (its manifest.npz is located inside, mirroring TreeDataSource) or an explicit manifest file path.

• New documentation guides: Balanced Datasets, Windowed Datasets, ArgBind Guide, Transform Chaining, Dict Batches, and Multiprocessing.

• New tests: Balanced datasets (40 tests covering hierarchical directory structures, unbalanced group sizes, statistical accuracy of weight-based sampling, and edge cases), multiprocessing/multithreading (11), transform chaining (8), and Dict[str, AudioTree] batch processing with scope (15).

• librosa dependency (>=0.11.0): decoding in AudioTree.from_file(), the CPU (soxr) backend for AudioTree.resample() and the windowed-LUFS cache, and NumPy-based STFT operations in the phase transforms (alongside librosax for the JAX side).

• Pre-commit + Ruff: Added a .pre-commit-config.yaml running Ruff (lint with --fix and Black-compatible formatting) plus basic file-hygiene hooks. pre-commit is included in the dev extra; enable with pre-commit install.

• CI covers macOS and Python 3.14: 3.14 was advertised in the classifiers but unreachable while bagz was a hard dependency; the macOS jobs additionally prove the no-bagz path.

• @random_transform / @map_transform are exported from audiotree.transforms and audiotree.transforms.jax. They are the documented way to write a custom transform, but lived only in the internal audiotree.transforms.decorators module — so the advertised extension point was not importable from any public namespace. They are now public API.

• README rewrite: the README was seven lines; it now carries the field table, install (including the bagz extra), a runnable quickstart, a comparison to Descript audiotools and torchaudio, and an explicit scope/non-goals list. The guides document the on-disk format rules (writer guide) and the reproducibility caveats (sources guide). The tagline “ML-framework-agnostic” is corrected in both the README and the PyPI description: JAX and Flax are hard requirements.

• on_read_error: create_audio_dataset, create_balanced_audio_dataset and AudioDataSource take on_read_error: "raise" | "skip" | "warn", defaulting to "raise". One corrupt file in a 100k-file corpus used to end a multi-hour run with no way to survive it. The non-raising policies drop the item using grain’s own convention: the loader returns None, later .map() stages are never called on it, and to_iter_dataset() skips it, so batches refill with real audio and nothing synthetic enters the pipeline (under random access, ds[i] returns None for a broken file). New public names, exported from audiotree.sources: AudioReadError (an OSError carrying .file_path) and OnReadError. Every read failure now surfaces as AudioReadError with the original as __cause__, so callers no longer have to know that librosa’s fallback raises audioread.NoBackendError while soundfile raises a RuntimeError subclass.

• TreeWriter(on_overflow=...) and a neural-codec guide (docs/source/introduction/codecs.rst) — AudioCodec is public API and appeared in zero guides. Writing it surfaced that LatentAudioCodec had never been exported, so every :class: role pointing at it was dead; both protocols are now exported from audiotree.transforms and audiotree.transforms.jax.

• tests/test_seed_stability.py pins the exact (basename, offset) sequence a fixed seed produces, and the multiprocessing tests now compare mp_prefetch against the single-process stream. Shuffling, excerpt selection and seed derivation were all edited in this release and nothing would have noticed a change to any of them.

• tests/transforms/test_backend_parity.py compares every transform exported by both backends, over mono and stereo. Every prior JAX transform test was mono, which is exactly why the shift_phase divergence below went unseen.

Changed

• Breaking — codec transforms take a codec object: encode_with_codec(codec) and encode_latents(codec) now take a single codec object instead of a bare encoder_fn (and, for encode_with_codec, a num_codebooks count). Two structural protocols in audiotree.transforms describe them: AudioCodec, whose encode(AudioTree) -> codes produces discrete tokens, and LatentAudioCodec, whose encode_to_latent(AudioTree) -> latents produces continuous latents. They are deliberately separate and both runtime_checkable, so an encode-only codec is expressible, a codec that does both simply satisfies both, and isinstance(codec, AudioCodec) is a real conformance check. The codec owns resampling, channel handling, and output shapes, so the transform no longer rearranges codes into a (batch, codebooks*channels, frames) packing: AudioTree.codes stores exactly what the codec returns; any state a codec keeps beyond the codes (such as a loudness-normalization factor its decoder needs back) is the codec’s own to manage. Both transforms remain idempotent (an AudioTree that already has codes/latents passes through) and are now defined once and shared by the audiotree.transforms and audiotree.transforms.jax namespaces.

• Breaking — Renamed the loudness vocabulary to lufs: The AudioTree.loudness field is now AudioTree.lufs, and AudioTree.replace_loudness() is now replace_lufs(). This affects attribute access and all keyword arguments (AudioTree(lufs=...), create(lufs=...), from_file(lufs=...), tree.replace(lufs=...)). The excerpt config’s loudness_cutoff is now ExcerptConfig.lufs_cutoff (see the SaliencyParams → ExcerptConfig rewrite below). Every new 1.0 surface follows the same spelling: normalize_lufs(), clear_lufs(), keep_lufs=, filter_by_lufs(), the lufs manifest columns and leaves, and the windowed-LUFS cache.

• Breaking — Renamed AudioTree.audio_data to AudioTree.waveform: The primary audio field is now called waveform (matching torchaudio’s convention and avoiding ambiguity with sample counts). This affects all keyword arguments (AudioTree(waveform=...), tree.replace(waveform=...)) and attribute access. It also changes the on-disk leaf names produced by TreeWriter: new datasets are written with waveform.bin (or e.g., dry.waveform.bin when nested) and matching manifest.json keys. To migrate an existing dataset, rename each *audio_data.bin file to the corresponding *waveform.bin and replace every occurrence of audio_data with waveform in its manifest.json.

• Breaking — Separate NumPy and JAX transform backends: Transforms now have two implementations:

  • audiotree.transforms (NumPy): For CPU-based grain data pipelines. Uses np.random.Generator.

  • audiotree.transforms.jax (new): For GPU/JIT training pipelines. Uses jax.random.key.

  The base transforms in audiotree.transforms now expect np.random.Generator instead of jax.Array for the RNG parameter. Use np.random.default_rng(seed) instead of jax.random.key(seed). For JAX-based transforms in jitted training loops, import from audiotree.transforms.jax.

• Breaking — Transform API redesign: Completely redesigned from class-based to function-based pattern. All transforms are now functions decorated with @random_transform or @map_transform, enabling direct parameter binding with argbind. Parameters are now flat (e.g., volume_norm.min_db: -20) instead of nested in a config dict. This allows CLI parameter binding (e.g., --volume_norm.min_db=-20) which was not possible with the old API. All transform names changed from PascalCase to snake_case (e.g., VolumeNorm → volume_norm, Trim → trim). Old class-based transforms removed with no backward compatibility. Migrated transforms: identity(), mono(), stereo(), volume_change(), volume_norm(), rescale_audio(), invert_phase(), swap_stereo(), corrupt_phase(), shift_phase(), roll(), trim(), choose(), encode_with_codec(), encode_latents(). The reserved parameters (prob, split_seed, scope, output_key) are keyword-only, and transform instances are picklable with a real __repr__ (grain spawns its workers rather than forking).

• Breaking — Renamed module audiotree.datasources to audiotree.sources: To match grain.sources.

• Breaking — Refactored create_balanced_audio_dataset(): Now uses grain’s random_map pattern for saliency-based audio loading (see Fixed for the excerpt-diversity bug this resolves). Now supports mixing pre-constructed grain MapDatasets via the datasets parameter. The sources parameter is now optional (at least one of sources or datasets must be provided). Pre-constructed datasets passed via datasets must already be repeated (call .repeat() before passing); if a finite dataset is passed, grain.MapDataset.mix will truncate output to the shortest dataset length. File-based sources are automatically repeated internally.

• Breaking — Renamed seed to shuffle_seed and added excerpt_seed: In create_audio_dataset() and create_balanced_audio_dataset(), the seed parameter has been replaced with shuffle_seed (controls file order) and excerpt_seed (controls random excerpt selection). This allows creating datasets that visit files in the same order but load different random excerpts. excerpt_seed=None still falls back to shuffle_seed, but the two streams are drawn as separate children of one numpy.random.SeedSequence rather than being the same integer, so shuffle order and excerpt offsets stay decorrelated even under the fallback. In create_balanced_audio_dataset(), each group’s seeds are derived from the group’s name — SeedSequence([base_seed, crc32(name), crc32(role)]) — so adding or reordering a group leaves every other group’s stream untouched, and the shuffle and excerpt roles stay independent even when the two base seeds are equal.

• Level-, rate- and channel-changing AudioTree methods work on mini-batched trees: After reshape_mini_batches(), the waveform has shape (num_mini_batches, mini_batch_size, C, T). replace_lufs(), normalize_lufs(), to_mono(), to_stereo(), and resample() now handle any number of leading batch axes by flattening them around the underlying 3-D kernels and restoring them afterwards (lufs comes back shaped like the leading axes, e.g., (num_mini_batches, mini_batch_size)). The methods whose meaning is genuinely per-item still require rank 3 and now say so by name — AudioTree.write() and AudioTree.create() raise naming the rank and the fix rather than silently treating the mini-batch axis as the batch.

• AudioTree array fields typed as ArrayLike: The waveform, lufs, lufs_windows, pitch, velocity, note_duration, codes, and latents fields (and the matching create() / from_file() parameters) are now annotated ArrayLike = np.ndarray | jax.Array instead of np.ndarray, so type checkers accept both NumPy and JAX arrays. Optional fields are explicitly ArrayLike | None, and AudioTree.create() now accepts waveform=None (the dimensionality-normalizing reshape is skipped) for token-only trees built from codes / latents. core.py uses from __future__ import annotations, so no annotation is evaluated at import time, and ArrayLike widens to the union only under TYPE_CHECKING (it is plain np.ndarray at runtime). Runtime behavior is otherwise unchanged.

• Enhanced loudness computation: AudioTree.replace_lufs() now uses the loudness library for numpy arrays and jaxloudnorm for JAX arrays.

• AudioTree.resample() CPU backend: NumPy-backed waveforms now resample on CPU via librosa (soxr), and the result stays a NumPy array — previously a NumPy waveform was silently returned as a JAX array. JAX-backed waveforms still use the JAX/Julius port. The NumPy path’s output length is pinned to match the JAX backend (soxr can otherwise differ by a sample). zeros, rolloff, and full apply to the JAX backend only.

• jaxloudnorm dependency: Now installed from PyPI and pinned to >=0.3.2. The 0.3 line returns -inf LUFS for digital silence (the mathematical limit of zero gated power), so AudioTree.replace_lufs() reports a well-defined value on silent JAX audio instead of NaN; the floor is .2 rather than .1 because jaxloudnorm<=0.3.1 depended on the 2017 pysoundfile distribution, which installs its own top-level soundfile.py and silently shadowed the real soundfile in the lowest-version resolution.

• Loosened the grain pin to >=0.2.15,<0.3 (was ==0.2.15): nothing in the package rides grain’s private API any more. The one class that did, ReduceBatchTransform, which subclassed a grain._src batch operation, is gone (see Removed); batching is now AudioTree.batch over jax.tree_util.tree_map. Both ends of the range are exercised in CI: the test_lowest_direct job resolves 0.2.15 and the normal matrix takes the newest 0.2.x. Note: grain 0.2.17+ reads absl flags inside its multiprocessing prefetch, so scripts doing multiprocessing data loading outside an absl.app.run entry point must call flags.FLAGS.mark_as_parsed() once at startup (the test suite does this in tests/conftest.py).

• Filepath/source strings raise instead of truncating: Encoding a filepath or source string longer than audiotree.core._str_max_length now raises ValueError instead of silently truncating it (which corrupted the round-tripped path). The default limit was also raised from 256 to 1024 characters.

• Updated examples: argbind_augmentations examples demonstrate new function-based transform API with simpler YAML configs and working CLI parameter binding.

• Consolidated tests: Merged test_roll.py into test_core.py for better organization.

• Breaking — bagz is now an optional extra: install audiotree[bagz] for string leaves in TreeWriter/TreeDataSource and for the windowed-LUFS cache. It was a hard dependency gated on sys_platform == 'linux', but bagz publishes manylinux x86-64 wheels only and has no sdist, which made pip install audiotree unsatisfiable on Linux aarch64 at every Python version and on Linux x86-64 at Python 3.14. Relatedly, exclude_prefixes now actually escapes bagz: the import was previously required before the exclusion filter, so excluding every string leaf still raised ImportError and a Linux-written dataset could not be opened elsewhere at all — not even to read its waveforms.

• jax, flax and absl-py are now declared dependencies: all three are imported at module scope but arrived only transitively via jaxloudnorm/librosax, so import audiotree depended on another package’s resolution choices.

• py.typed: the package now ships a PEP 561 marker, so its annotations are visible to downstream type checkers instead of being silently ignored.

• License metadata: declared via PEP 639 (license = "MIT" plus license-files), which replaces the now-forbidden license classifier. Third-party notices for the julius-derived resampler and the pyloudnorm-derived loudness code are bundled under LICENSES/ and ship in both the wheel and the sdist. Audio fixtures under tests/assets/ are explicitly carved out of the MIT grant — the MUSDB18-HQ excerpt is CC BY-NC-SA 4.0 and cannot be relicensed.

• Breaking — transform constructors reject unknown parameters: a misspelled name (volume_change(min_dB=40), capital B) was accepted and silently dropped, so the augmentation a config asked for simply never happened. It now raises TypeError with a suggestion. A config that appeared to work may now correctly fail.

• Breaking — scope={'key': True} raises: a scope entry is matched one level above where it sits, so a bool directly under a top-level key selected every leaf rather than that one — silently augmenting the target signal in a dry/wet pipeline. Use the new list form, scope=["wet"] (also ["input.dry"] or [("input", "dry")]); the nested-dict sentinel form {"wet": {"scope": True}} is unchanged. An explicitly empty scope ([] or {}) now selects nothing rather than everything — “transform these subtrees: (none)” makes the transform a no-op instead of silently applying it to every leaf.

• Breaking — public callables are keyword-only past their leading positional arguments: signatures exposed up to 16 positional parameters (AudioTree.from_file), freezing their order at 1.0. Most now take one or two positionally; create_balanced_audio_dataset is the widest at three (sources, weights, datasets). tests/test_public_api.py pins the per-callable budget, so adding, reordering or renaming a parameter cannot silently change what a positional call means.

• Breaking — AudioTree.filter’s callback is predicate, not filter_fn, and the method is documented (it previously had no docstring and so was missing from the API reference).

• Dependency floors are now the versions that actually work: verified by resolving every direct dependency to its lowest allowed version and running the suite there, which the new test_lowest_direct CI job repeats. numpy>=1.24, scipy>=1.10 and librosa>=0.10.1 were unreachable fiction (the graph forces 2.1.3 / 1.16.2 / 0.11.0); jax and flax needed raising to >=0.10.0 and >=0.12.3, since older releases break on the funless @jax.jit(...) form and on nnx.jit.

• Breaking — AudioTree.create/from_file take filepath=, not filepaths=: the kwarg was the odd one out three ways — the property is .filepath, the manifest column is filepath, and the sibling source= is singular while also accepting a per-item list. The rule is that a plural name means many values per item (codes, latents), not many items; filepath, source, pitch and velocity are one per item, and a list is just the batch axis. from_file returns a batch of 1, so it accepts a single path or a one-item list; a longer list used to be stored as-is, silently misaligning the one-row-per-item provenance the moment such trees were batched, and now raises. create_audio_dataset(filepaths=...), scan_durations(filepaths) and precompute_window_lufs(filepaths) keep the plural — those genuinely take many files.

• Breaking — SaliencyParams is now ExcerptConfig, built around a named strategy: the old object offered four spellings for three behaviors — saliency_params=None and enabled=False both meant “offset 0”, enabled=True, lufs_cutoff=None meant “random offset”, and enabled=True, lufs_cutoff=-40 meant “search for a loud section”. One strategy field now names them: "start", "random" (the default) and "loudest". This fixes a silent data bug: because the dataset constructors defaulted to None, create_audio_dataset(...) out of the box read the first duration seconds of every file on every epoch, and excerpt_seed controlled nothing — while the docstring example claimed otherwise. The default is now a random offset.

  • on_failure decides what happens when no candidate clears lufs_cutoff. "loudest" is a best-of-num_tries search, not a filter, so it previously returned the quietest excerpt it found with nothing to distinguish that from success. "keep" (default) preserves the old behavior, "skip" returns None (grain drops it at to_iter_dataset()), and "raise" raises naming the file.

  • Knobs that do nothing now raise. Setting lufs_cutoff, num_tries or on_failure under a strategy other than "loudest", or search under "start", was silent; it is now a ValueError at construction. (search is legal under "random", whose single offset it draws.)

  • It is a plain frozen dataclass, not a flax.struct.dataclass. It was registered as a JAX pytree whose leaves were [True, 8, -40.0, 'uniform'] — including a string leaf — despite being pure CPU-side config that is never traced.

  • search_function is now search, annotated str | Callable rather than lying with str, and the two built-ins are module-level search_uniform / search_bias_early instead of staticmethods on the config. AudioTree.salient_excerpt() is now AudioTree.loudest_excerpt(), taking excerpt=. ExcerptConfig is exported from audiotree.sources as well as audiotree — it is a sources concept and was missing from that namespace.

  • create_balanced_audio_dataset rejects window together with a customized excerpt (previously with a non-None saliency_params).

• WindowParams is WindowConfig, passed as window= (renamed after 1.0.0rc1, which shipped the old spellings): the class name matches ExcerptConfig, and the create_balanced_audio_dataset(window=...) keyword matches excerpt=.

• The merged metadata dict is split: payload is AudioTree.extras, provenance moves to the private AudioTree._metadata container (finalized after 1.0.0rc1, which shipped the old merged dict — spelled metadata, with user payload and the library’s bookkeeping side by side; no rc2 was ever published). extras is the user’s dict of per-item array/string leaves that batch with the waveform — labels, embeddings, features — and it is entirely the user’s: the library plants no keys of its own there; every API spelling follows (create(extras=...), from_file(extras=...), tree.extras[...], replace_extras()). _metadata is now a library-internal provenance container — private, as the underscore says, serialized on disk under the name metadata — holding exactly the encoded filepath and source arrays (fixed-width integer encodings that survive jit and batching) plus the plain offset array (where in its source file each excerpt starts): users pass filepath=/source=/offset= to create() (or let from_file and the loaders stamp them) and read the .filepath/.source/.offset properties, never the container itself, and its schema is closed — a metadata node read off disk may contain only filepath/source/offset, anything else is rejected by name, the same strict-validation stance as manifest columns. Each name is now honest — extras is payload that trains with the audio, the per-item provenance container is data about the audio (where it came from) — and neither collides with TreeWriter(metadata=...), the true dataset-level metadata (the manifest’s top-level "metadata" key, read back via TreeDataSource.get_metadata()), which keeps its name. On disk: a TreeWriter AudioTree node serializes both children (metadata with the filepath/source/offset leaves, extras with user leaves as extras.<key>), and an AudioWriter NPZ manifest stores filepath, source and offset as dedicated bookkeeping columns, leaving the extras_<key> columns (masks __mask_extras_<key>) purely user payload. There is no read-side alias — TreeDataSource rejects an unknown structure child by name, and the NPZ readers reject any unrecognized manifest column by name. For an rc1-era dataset the consequence is schema-shaped: a legacy metadata node holding only filepath/source reads correctly with its original meaning, while one holding user payload fails loudly at read and must be rewritten. The rewrite is mechanical: move payload leaf paths/files from metadata.<key> to extras.<key> in manifest.json (filepath/source stay put), and rename NPZ payload columns metadata_*/__mask_metadata_* to extras_*/__mask_extras_* and metadata_source to source (filepath was already a dedicated column).

• Breaking — the three on-disk formats are versioned: a TreeWriter directory, an AudioWriter NPZ manifest, and a windowed-LUFS cache each now carry a header (format, format_version, min_reader_version, producer), and every reader validates it. A major-version mismatch is refused; minor versions are additive so older readers keep working; min_reader_version is what a writer raises when it adds a field readers must honor. Pointing a reader at the wrong kind of directory now fails by name instead of as a KeyError. Previously TreeWriter stamped a "version": "2.0" that was never bumped for the audio_data → waveform rename, the NPZ manifest had no version field at all, and the LUFS cache wrote a version its reader never checked (a hand-edited "99.0" loaded happily). Artifacts written by a pre-1.0 audiotree carry no header and are refused with an explicit “re-render the dataset” message — none of these layouts shipped in a release (1.0 is the first to contain TreeWriter), so there is no released data to migrate.

• Breaking — swap_stereo documented “must be stereo” while no-opping on mono and reversing channel order for C>2. One behaviour now, identical in both backends.

• Derived fields invalidate in one place. resample, to_mono and to_stereo cleared lufs but left codes and latents stale — and encode_with_codec is idempotent, so resample-after-encode reused tokens describing the old audio. The transform layer had the same gap: volume_norm, volume_change, roll, trim, the phase transforms, rescale_audio and peak_norm all changed the waveform while preserving codec tokens for audio that no longer exists. One _invalidate_derived() helper now covers every length, rate, channel and energy changing operation on both surfaces; invert_phase and swap_stereo keep the (genuinely unchanged) loudness but drop codes/latents, which describe the un-negated samples and the un-swapped channel order.

• Public asserts are exceptions. They vanish under python -O, so every guarantee they enforced silently disappeared in a tuned training run.

• The two loudness engines agree closely, not exactly. replace_lufs() measures with the exact ITU-R BS.1770 IIR meter on a NumPy waveform and with jaxloudnorm’s FIR-approximated K-weighting on a JAX one, choosing from the waveform’s own array library unless you pass engine=. They are not bit-identical. Since lufs is persisted as a manifest column and drives lufs_cutoff corpus filtering, pin engine= explicitly when a corpus has to be reproducible across backends.

• Constructors and writers reject inconsistent input at the boundary, rather than accepting it and misbehaving later: AudioTree.create refuses a filepath/source list whose length does not match the batch (it would otherwise misalign provenance); create_balanced_audio_dataset refuses a group name present in both sources and datasets (it silently doubled that group’s share); create_windowed_audio_dataset requires a positive duration and, when loudness filtering is active, a positive lufs_window_sec (a non-positive value built a dataset of zero-length windows or filtered against a fabricated grid); AudioTree.create rejects a Python list/tuple of non-strings as an extras leaf (jax.tree_util descends into such a list, so indexing and filter passed it through at full length, silently misaligned with the batch — convert with np.asarray; lists of strings remain a supported per-item leaf); roll requires min_seconds <= max_seconds on both backends (NumPy raised, JAX silently rolled every item by min_seconds); and AudioTree.from_file raises when a nonzero offset reads past the end of the file (padding the empty read out used to return full-duration silence stamped with an offset pointing at audio that does not exist — silent zero training data for any caller with an off-by-one, or seconds/samples confusion, in its own offset math); and AudioWriter requires an {index} field in the filename pattern when writing audio (without it every item overwrote one file while the manifest recorded a row per lost item) and validates the manifest-column schema and filepath coverage on every write, so a drifting write raises at its own call — before its audio lands — instead of aborting at manifest save with the manifest unwritten.

• TreeWriter.write is atomic and dtype-checked: it validates every leaf (batch size, per-sample shape, exact dtype, and string encodability) before touching any file, so a bad batch is rejected whole with nothing written, and a leaf whose dtype differs from the schema fixed on the first write now raises instead of being unsafe-cast into the memmap (an int16 schema silently storing 1e9 read back as -13824). A failure during the commit poisons the writer so close() cannot finalize a shifted dataset. Initialization is held to the same standard: the schema is recorded only after every leaf file is created and mapped (a failure creating a later leaf’s file used to leave more leaf names than memmaps, and the commit loop then silently dropped the trailing leaves of every subsequent write while reporting success — the files created so far are now removed and the writer poisoned); a 0-d array leaf is rejected by name (every array leaf needs a leading batch dimension) instead of dying in a bare IndexError after files exist; and the constructor verifies metadata= is JSON-serializable up front, since it only meets json.dump at manifest time — after the leaf data has committed — where a stray numpy scalar left an hours-long render with all its data and no manifest at all.

• The manifest format stores one type per column and refuses values it cannot round-trip: a column mixing logical types (e.g. a bool column containing an int, or an int column containing a float) now raises rather than coercing (5 → True, 2.5 → 2); a string or bytes value ending in a NUL raises (fixed-width <U/|S storage drops a trailing NUL, breaking the round-trip contract); and the column name tags is reserved, since read_entries synthesizes it from the tags_* columns.

• create_balanced_audio_dataset stamps the full provenance schema onto pre-built datasets: source= is set to the group name (overwriting any source they carried), a missing filepath is filled with empty strings, and a missing offset with NaN — so an item built with AudioTree.create (which stamps none of the three) exposes the same per-item metadata schema as the file-based groups, and a batch spanning both no longer fails its pytree check on the missing keys — a failure that was intermittent, since single-group batches collated fine. The function’s own docstring example mixes the two. When channels is left to auto-probe (and mono=False), one file is probed and the count applied to every group, so two internally-consistent groups that disagree with each other (stereo music, mono speech) fail at load time with a filename instead of at batch time with a shape error that names nothing.

• Dataset seeds use their full width: shuffle_seed/excerpt_seed were masked to their low 32 bits before seeding numpy.random.SeedSequence, so seeds differing only above bit 31 (e.g. 1 and 2**32 + 1, or a 64-bit hash / time.time_ns()) produced byte-identical datasets. The full seed is now used; small-seed outputs are unchanged (masking a sub-2**32 seed was a no-op), so existing runs reproduce.

• TreeDataSource validates the whole manifest at construction, by name: beyond the format header and traversal checks, it now validates string_leaves entries (file type, in-directory path, existence), uses arbitrary-precision arithmetic for the leaf size check (an overflowing shape_per_sample could wrap it to zero and pass), rejects a structure child colliding with a reader-set field such as sample_rate, and requires the top-level keys — each a named ValueError at construction instead of an unnamed TypeError/KeyError deep in a grain worker. It also cross-checks container types and reference integrity: "leaves": [] or a dict node without children used to crash with a raw AttributeError/KeyError at first read, and a structure node referencing a leaf that exists in no leaf table silently dropped the field (a dangling root even leaked an internal sentinel object) — all are now named “Invalid manifest” errors at construction. The windowed-LUFS cache reader checks its required manifest keys the same way, before the bagz file is opened.

• A string output_key requires exactly one in-scope leaf, and scope entries reject unknown keys: with several in-scope leaves, a literal output_key mapped every result to the same name, so all but the dict-order-last transformed subtree was computed and silently discarded — costly with a codec transform; it now raises naming the colliding keys (use a callable output_key for several leaves). Non-AudioTree entries of a dict element (label arrays, ids) no longer participate in the rename at all: they used to count toward the single-leaf rule — so a string output_key on {"audio": tree, "labels": array} reported a spurious collision despite there being exactly one transformable leaf — and a callable output_key silently duplicated each of them under the new name. A string output_key also no longer costs picklability (it became a lambda defined inside __init__, which stdlib pickle refuses — quietly contradicting the picklable-transforms rule grain’s spawned workers rely on). A scope dict entry carrying keys other than the scope sentinel — the shape one page of the docs wrongly taught as per-key parameter overrides — was silently misread as extra scope markers; unknown keys now raise, and the docs teach the real pattern (one scoped transform instance per key).

• Parameter validation happens at construction and identically on both backends: an inverted min/max range in volume_change/volume_norm/corrupt_phase/shift_phase made NumPy raise numpy’s generic error per element while JAX silently clamped every draw to the minimum (a constant gain, no error); a clip shorter than frame_length made librosa quietly shrink the FFT while JAX crashed; and choose() accepted weights that were negative or did not sum to 1, failing minutes later inside a grain worker. All three now raise the same named ValueError on both backends, at construction where possible. Map-transform error messages stop advertising prob/split_seed (which they reject), and a transform’s repr includes non-default prob/split_seed/scope/output_key instead of rendering a constructor call that would rebuild a different transform.

• Breaking — corrupt_phase/shift_phase require hop_factor in (0, 0.5]: above 0.5 the hann analysis windows no longer overlap-add to a constant, so even amount=0.0 — mathematically the identity — corrupted chunks of every frame at hop_factor=1.0, and between 0.5 and 1.0 the NumPy and JAX backends silently disagreed on the tail samples (NumPy left the last partial hop unanalyzed; JAX padded and reconstructed it). A hop_factor small enough to truncate to a zero-sample hop failed differently per backend (librosa ParameterError vs. a bare ZeroDivisionError). Both cases now raise the same named ValueError on both backends, like the other transform parameter guards.

• volume_norm trusts a cached lufs: it re-measured every element unconditionally, even when the tree arrived with lufs populated — for instance restored from an AudioWriter manifest, whose whole point is to make re-measurement unnecessary — running the one-item-at-a-time CPU meter on every element of every epoch. It now measures only when lufs is unset, the same rule normalize_lufs and volume_change already follow; the cache is trustworthy because every operation that changes the audio clears it.

• AudioWriter pins each manifest column’s kind at its first write: a column whose values drifted in logical kind across writes (int rows, then a str row) passed every write() and aborted only at save/close — manifest unwritten, every WAV orphaned. The offending write() now raises, naming the column and both kinds, before its audio lands; close-time validation stays as the backstop.

• create_balanced_audio_dataset rejects non-positive group weights by name: 0, negative, NaN and inf weights previously fell through to grain’s mixer, which failed with errors naming neither the group nor the argument; omission, not zero-weighting, is the way to drop a group.

• AudioTree.excerpt requires duration honestly: the parameter was annotated Optional[float] = None and documented as optional while unconditionally raising on None; it is now a required argument.

Removed

• AudioDataSimpleSource, AudioDataBalancedSource, AudioDataBalancedDataset, and AudioDataSourceMixin: These deprecated classes are removed. Use create_audio_dataset() in place of AudioDataSimpleSource, and create_balanced_audio_dataset() in place of the other three.

• AudioTree.from_array(): Use AudioTree.create() instead.

• num_records parameter: Removed from create_audio_dataset() and create_balanced_audio_dataset(). Use .slice(slice(0, N)) on the returned dataset instead. This simplifies the API and follows the principle of separation of concerns.

• Post-mix shuffle in create_balanced_audio_dataset(): The automatic shuffle after mixing datasets has been removed. The shuffle parameter now only controls whether files within each group are shuffled. If you need the mixed output shuffled, call .shuffle(seed=N) on the result. This change simplifies the API and gives users explicit control.

• Breaking — ReduceBatchTransform is removed: it subclassed a grain._src class grain logs as deprecated and overrode three private members. Use ds.to_iter_dataset().batch(n, batch_fn=AudioTree.batch) instead; the new AudioTree.batch collation preserves the array library (JAX in, JAX out), so batching a device-resident tree is not a blocking host sync.

Fixed

• The windowed-LUFS reduction no longer materializes every window: an 84.7 MB batch at the standard EBU setting (3 s window, 100 ms hop) allocated ~29× its input. It is a sliding sum, so it never needed the gather — sliding_window_view on NumPy, lax.reduce_window on JAX. Peak RSS above baseline falls from 2451 MB to 338 MB (28.9× → 4.0× of input) and the reduction is 30× faster. It is also slightly more accurate than the gather it replaces: each window is summed over contiguous memory in ascending order, which reproduces a per-window mean() exactly and sits 1.5e-06 from a float64 reference, where reducing the materialized (items, windows, span) copy blocks differently and drifts to 5e-05 on a multi-item batch (the two agree bit for bit only for a single mono item). JAX sits within 4.1e-06 relative of the NumPy path, i.e. float32 rounding rather than reordering.

• JAX shift_phase decorrelated stereo: it drew one phase angle per (batch, channel) where NumPy drew one per item, so a stereo pair whose channels start identical came apart — L-R maxdiff 0.86, negatively correlated in half of the draws sampled. Both backends now rotate per item. corrupt_phase still draws per channel on both, so the two transforms remain distinct.

• JAX corrupt_phase/shift_phase zeroed the signal tail whenever length % hop_length != 0 — 68 samples on a 1 s @ 44.1 kHz clip, an audible discontinuity. Corrected locally by padding to a hop multiple and trimming after.

• Mini-batched (rank-4) trees wrote nonsense instead of failing. AudioWriter.write read the mini-batch axis as the batch, shifting every axis below it: a 6-item (3, 2, 1, 800) tree produced three manifest rows claiming channels=2, samples=1. write_audio=True was saved only by soundfile rejecting the shape; the manifest-only mode recorded it silently. .filepath/.source now decode per mini-batch rather than raising, and the operations with no rank-4 reading say which rank they got.

• TreeDataSource validates leaf files against the manifest. A manifest claiming more samples than the .bin holds surfaced as a bare np.memmap error naming neither the manifest nor the leaf — and where memmap tolerated it, not at all. The bound is >=, not ==, because flush() is public and a mid-write .bin is legitimately longer.

• load_into_memory=True handed out writable aliases into the shared store, and returned empty samples in grain worker processes when the parent had read first.

• Breaking — two transform guards were asserts, so python -O removed them: BaseRandomTransform.__init__ checked 0 <= prob <= 1, and BaseTransformMixIn._post_process checked that a transform with an output_key was handed a dict, both with a bare assert. Under -O the first let prob=1.5 through as “always” (and prob=-1 as “never”), and the second let the rename path run on a non-dict element. Both now raise — ValueError matching choose()’s wording, and TypeError naming the type it actually got.

• Breaking — TreeWriter wrote three things TreeDataSource refuses: a non-native-endian leaf recorded ">f4" against a reader allow-list that accepts only native names; two leaves whose dot-joined paths collide ({"a.b": x} and {"a": {"b": y}}) shared one a.b.bin, so y vanished and both read back as x; and a leaf named "../escaped" wrote outside the dataset directory. Endianness is now normalized on write, colliding filenames raise naming both leaves, every write-side path goes through the same safe_join guard as the reader, and a dtype the reader would reject fails on the first write rather than after an hours-long render.

• Breaking — the loudness meter measured padding as if it were audio: excerpts shorter than one 400 ms BS.1770 gating block were zero-padded, diluting mean square by dur / 0.4 and under-reporting by exactly 10*log10(dur/0.4) — a 0.2 s tone measured −6.05 LUFS where the truth is −3.01. The padded signal is now level-compensated before the meter, so the absolute and relative gates see the corrected level and the reading is duration-independent. Separately, the JAX FIR K-weighting was pinned at 512 taps regardless of sample rate and so could not realize the 38 Hz high-pass much above 16 kHz (agreement with the exact IIR path at 96 kHz improves from 1.52 dB to 0.008 dB), and lufs and lufs_windows were computed with different K-weighting filters within one NumPy call — BS.1770 DeMan for the integrated path, pyloudnorm’s legacy RBJ approximation for the windowed one. Both now use BS.1770. These are corpus-affecting: lufs is a persisted manifest column and the input to lufs_cutoff filtering.

• Breaking — split_seed=False meant two different things per backend: JAX shared an immutable key, so every leaf drew identically; NumPy shared a stateful Generator, so each leaf advanced the stream. A dry/wet pair identical on JAX came out 8.8 dB apart on NumPy — the grain/CPU default — silently training effect-modelling pipelines on decorrelated pairs.

• Breaking — filter_by_tag and filter_by_lufs discarded the receiver’s filter: both return an AudioDataSource and read as a fluent API, but neither filter_fn nor num_records was carried on self, so each helper re-read the whole manifest from disk. Chaining filter_by_lufs onto a filter_by_tag result brought back exactly the records the user had excluded, and a constructor-supplied filter_fn was dropped by a single helper call. Filters now compose.

• Breaking — per-group RNG seeds were bound to dict position: inserting or reordering a group in sources swapped entire per-group streams, silently changing what every downstream run trained on. Seeds now derive from the group name.

• AudioTree.excerpt(search_function=) could not call either shipped searcher: it was invoked with four positional arguments while search_uniform, search_bias_early, and everything _resolve_search_function returns need six, and passing a registered name failed with 'str' object is not callable. Two incompatible searcher protocols lived in one module; excerpt now takes an ExcerptConfig like loudest_excerpt, leaving one signature and one resolver.

• t[np.int64(i)] silently dropped the batch axis: __getitem__ special-cased only the builtin int, so t[np.argmax(t.lufs)] — or any loop over np.arange(n) — scalar-indexed every leaf and returned a tree whose batch_size reported the channel count. On a waveform-only tree nothing raised at all. Integer scalars of any array library now slice correctly; bool is excluded so t[True] no longer selects index 1. A Python list or boolean-mask key is normalized to an array once, up front, so it indexes a JAX-backed tree the same way it indexes a NumPy-backed one (JAX arrays reject a bare list key).

• create(filepath="a.wav") on a batch tagged only row 0 despite the docstring promising it tags the whole batch, so t[2].filepath was empty and any filter dropping item 0 lost provenance entirely — which then landed in written manifests. Same for source=.

• prob < 1 was unusable in two more configurations: combined with output_key it raised a pytree structure error on every seed while the constructor merely warned, and on the NumPy backend at batch_size == 1 it made the per-element pytree structure random, so AudioTree.batch died downstream — the README’s canonical random_map → batch(batch_fn=AudioTree.batch) shape.

• Dict, list, and tuple transform parameters: dict-valued parameters were silently replaced by their default (_get_config_val compared a dict’s inner key against the parameter name, and leaked that key into the sibling-parameter namespace); list and tuple parameters crashed outright.

• Decorated transforms could not have required parameters, and reported an “unexpected parameter” that the decorator’s own __signature__ advertised. The common prob/split_seed/scope/output_key documentation was appended after the Example block, so none of the 28 transforms documented any of it. Each now has exactly the reserved parameters it accepts spliced into its own Args: block — all four for a random transform, scope/output_key for a map transform, which is also the only pair a map transform will accept.

• resample crashed on bf16/fp16, and on fp32 under jax_enable_x64 — the sinc kernel took JAX’s global default dtype rather than the input’s. The equal-rate short-circuit also skipped every argument check, accepting output_length=-5 and full=True with output_length only for the files that happened to already be at the target rate. The convolution now runs at HIGHEST precision and the kernel is cached.

• Silently-ignored arguments now raise: unknown weights keys, duration combined with window, and an empty sources mapping (which failed with min() iterable argument is empty). A source matching zero files warns instead of shrinking the corpus in silence, pathlib.Path is accepted wherever a str path is, and a mixed-channel corpus names the offending file and its channel count rather than failing opaquely at batch time.

• show_progress=True was a silent no-op: tqdm was imported but declared in no dependency, extra, or group, and the ImportError was swallowed. It now raises with an install hint, and there is a progress extra.

• Every guide snippet was executable fiction: the 1.0 renames landed in the source but not in the guides, which use code-block:: python rather than doctests — so the docs build stayed green while three guides constructed ExcerptConfig(enabled=True) and two passed a seed= that had been split in two. tests/test_docs_snippets.py now extracts all 154 Python blocks under docs/source/: every block must parse, and every call to a known audiotree callable is matched against the live inspect.signature, so a renamed keyword fails even in a snippet that could never run for want of a corpus. Blocks that can run are then executed against a synthetic corpus built from the placeholder paths the guides themselves name; one that legitimately cannot opts out with a .. skip-snippet-exec: comment and a reason, and stays subject to every other check.

• The examples/ scripts run again, and CI runs them: all three were broken against the renamed loudness API (replace_loudness(), .loudness, filter_by_loudness), so a user’s first contact with the library tracebacked immediately, and testpaths excluded examples/ so nothing went red. writer_datasource_example.py additionally wrote one item without a pitch field, which AudioWriter correctly rejects.

• Silence no longer normalizes to an all-NaN waveform: an item whose measured loudness is not finite has no gain that reaches a target (target - (-inf) is +inf, and 0 * inf is NaN). AudioTree.normalize_lufs() and the volume_norm transform (both backends) scaled by that gain and then stamped the target LUFS on the result, so a silent item became an all-NaN waveform advertising itself as, say, −18 LUFS. This was never limited to digital silence: anything below the BS.1770 absolute gate reads -inf, which zero-padded short reads and trim(mode="constant") produce routinely, and one such item silently poisons every batch it lands in. Non-finite items are now passed through unscaled and keep their -inf, so they stay identifiable. normalize_lufs() also gains an opt-in max_gain_db ceiling so a very quiet but still measurable item is not amplified without bound.

• prob is drawn per batch item, and works on the JAX backend: random.bernoulli(key, p=prob) was called with no shape, so one scalar draw was broadcast across the whole batch — prob=0.5 on a batch of 32 augmented all 32 items or none. The marginal rate was correct, so this was invisible in aggregate while removing exactly the within-batch diversity prob exists to provide. Masking also went through jax.tree.map over the transformed and original trees, which requires identical treedefs, so any transform that nulls lufs (shift_phase, corrupt_phase, roll) or adds an extras key crashed outright at prob < 1, as did every transform on a tree carrying the string-encoded filepath provenance that from_file always sets — making prob < 1 unusable in the JAX namespace on anything from the data loader. Selection is now per AudioTree leaf with a (B,) mask, canonicalizing a field that only one side populates so the output structure never depends on the draw. A transform that changes a field’s shape cannot be mixed item-by-item and now raises a clear error. The prob key is also reserved before the per-leaf keys are derived; previously, for a two-leaf tree, the Bernoulli draw reused a leaf’s transform key.

• AudioDataSource no longer corrupts array fields, dtypes, or -1/NaN values: array-valued fields lost their batch axis, because AudioTree.from_file only adds one to a scalar — lufs_windows came back (W,) and codes (codebooks, frames), so AudioTree.batch concatenated along the wrong axis and interleaved one item’s tokens into the next (a one-window lufs_windows degraded to 0-d and made batching raise). Stored dtypes were re-cast, so an int32 pitch read back float32 and an int32 velocity read back int16. And -1, NaN, and "" were treated as “missing” and dropped, though nothing in AudioWriter writes them as sentinels: an item with velocity=-1 came back with a different pytree structure than its siblings and broke batching for the whole dataset. Every AudioTree field now gets an explicit batch axis on read and keeps its stored array and dtype. The derived fields (lufs, lufs_windows, codes, latents) are restored only when the source’s sample_rate/mono/duration leave the audio as it was written — they describe the written audio, and restoring them onto a resampled, mono-folded or trimmed load silently attached measurements of audio the caller was not getting. Annotation fields (pitch, velocity, note_duration) are restored either way.

• TreeWriter validates leaf names on every write: write() checked only the leaf count and then zipped positionally, while leaf extraction walks dict keys in sorted order. Writing {"a", "c"} after {"a", "b"} therefore stored c’s data in b.bin and the reader returned it under the key "b", with no error anywhere. sample_rate had the same problem — captured once from the first write, so a later tree at a different rate was permanently mislabelled (AudioWriter already raised on exactly that, so the two writers disagreed). The structure is now re-derived per write and must match, and the error names the differing leaves.

• TreeWriter publishes manifest.json eagerly: the manifest was written only at the bottom of close(), and flush() flushed the memmaps but not the manifest, so a multi-hour pre-render killed at 99% left valid .bin files that TreeDataSource refused to open at all. The manifest is now written as soon as the schema is known and refreshed from flush(), via a temp file plus os.replace so a concurrent reader never sees a partial one. A crashed render leaves a readable prefix.

• roll() invalidates the stale recorded offset: AudioTree.from_file records the source-file time of sample 0 as the offset provenance (read via .offset). Since roll shifts the waveform along the time axis (in both "wrap" and "constant" modes), that offset no longer points at sample 0 and is now set to None — mirroring how roll already invalidates lufs_windows in both modes, and lufs in "constant". The invalidation survives prob < 1: the per-item blend used to carry the original metadata through untouched, quietly re-attaching the stale offset to exactly the items that were rolled. Trees without a recorded offset are unaffected. Applies to both the NumPy and JAX backends.

• Stale loudness after length/channel changes: trim(), roll(mode="constant"), and AudioTree.to_stereo() (mono→stereo) now invalidate the cached lufs field, since changing the audio length, zero-padding, or duplicating a channel all change the integrated loudness. roll(mode="wrap"), invert_phase(), and swap_stereo() continue to preserve lufs because they leave it unchanged.

• Stale loudness after phase transforms: corrupt_phase() and shift_phase() now invalidate the cached lufs field by default. Their phase changes leave the magnitude spectrum (and thus energy) intact, so loudness is approximately unchanged, but the cached value is dropped to be safe. Pass keep_lufs=True to retain it.

• Excerpt diversity in create_balanced_audio_dataset(): record_key was used as both an array index and an RNG seed, severely limiting excerpt diversity when datasets were repeated. Fixed by the random_map refactor (see Changed).

• Manifest readers restore the source filepath: AudioWriter records each item’s source path as a top-level filepath manifest column (not under an extras_ prefix), but both readers previously dropped it. AudioTree.from_manifest() ignored the column entirely (so loaded.filepath came back empty even though the paths were on disk), and AudioDataSource reported the on-disk output path (or nothing, for manifest-only datasets) rather than the recorded source path. Both now restore the recorded source path into the filepath provenance, so .filepath matches what was written and the two readers agree; from_manifest keeps it aligned with the selected rows when a filter_fn is used. Manifests written without source paths are unaffected — from_manifest leaves .filepath empty, and AudioDataSource still falls back to the output audio path.

• output_key transforms on dict elements: a transform built with scope=[...] plus output_key=... raised IndexError on the first batch whenever the scoped-out key sorted alphabetically first (jax rebuilds dicts in sorted-key order, so {"dry", "wet"} with scope=["wet"] failed deterministically on both backends), and when the output_key collided with an existing key the merge silently kept the original instead of the transformed value (a pipeline regenerating "wet" each epoch trained on un-augmented audio). Both are fixed; the transformed value now wins for the key the user asked to write.

• AudioTree.split/filter/batch preserve list-of-strings extras and work on token-only trees: split sliced string-list extras (a supported leaf) character-wise rather than element-wise (["a.wav", "b.wav"] became ["a.", "b."]), and filter/batch then crashed on such trees; separately, these batch-axis operations dereferenced waveform.shape and raised AttributeError on a token-only tree (waveform=None, data in codes/latents). Both now read the batch axis from whichever leaf is present and treat a string list as one leaf.

• AudioTree.from_file honors duration without a target rate: target_length was computed only when both duration and sample_rate were given, so from_file(path, duration=...) at the file’s native rate returned a short file unpadded — despite the docstring’s unconditional promise — producing variable-length trees that crashed at batch time far from the cause. It now pads/trims to duration at the effective rate.

• Manifest round-trips array-valued and large/precise scalar columns: an array-valued string column came back as the literal repr string "['alice' 'bob']" (bytes as a raw fixed-width buffer) because the decoder stringified sub-arrays; and Python int/float tags were stored as int32/float32, so an int >= 2**31 raised OverflowError at save time (losing the whole manifest) and a float64 value silently lost precision. Array cells are now returned intact, and bare Python scalars widen to int64/float64. A scalar column’s dtype is the common type of all its rows, not the first row’s: pinning the first width silently narrowed later wider floats (out-of-range values became inf with only a numpy warning) and raised OverflowError on later wider ints — at close(), with every audio file already on disk and the manifest unwritten. int64 next to uint64, whose only common numpy type is a value-rounding float64, is refused by name.

• codec transforms are picklable: encode_with_codec / encode_latents returned instances of closure-minted classes that could not be pickled, so they failed at spawn time in grain multiprocess workers — unlike every other transform. They are now module-level classes holding the codec as an attribute; behavior is unchanged.

• TreeWriter.close no longer bricks a zero-size leaf: a leaf with a zero-size-per-sample shape (e.g. (batch, 0)) was truncated to 0 bytes at close, which cannot be memory-mapped, so the entire dataset raised “cannot mmap an empty file” at first access. Truncation is now floored at one byte, matching allocation.

• save_window_lufs writes its cache atomically: the bagz writer was closed with a bare call and the file written in place before the manifest, so a mid-write failure leaked the writer and could pair an old manifest.json with a partial lufs.bagz. It now closes in try/finally and writes to a temp path then os.replace, matching the manifest’s atomic write.

• String extras leaves survive every batch-axis operation: a bare str leaf (the batch-of-1 form TreeWriter documents and TreeDataSource produces) crashed AudioTree.batch — including the batch docstring’s own grain recipe — and List[str] leaves broke t[[0, 2]]/t[mask] indexing, were retained in full by a keep-nothing filter(), and stayed flat under reshape_mini_batches so indexing a mini-batch misattributed provenance. Every batch-axis operation now recognizes both string-leaf forms through one shared helper trio, and a torture sweep of both forms through every operation (now a test module) surfaced and fixed three sibling sites: strings next to trees in batch() were zipped into nested garbage, an emptied tree could not be re-filtered or re-batched, and a string leaf whose length disagreed with the batch slipped through reshape_mini_batches unchecked. filter/split on an empty tree also no longer raise ZeroDivisionError, and split(0) says what is wrong.

• The NumPy loudness engine enforces the documented 5-channel limit: replace_lufs promised a ValueError above five channels and the JAX engine raised one, but the NumPy engine silently computed a value; the check now runs before engine dispatch. A lufs_hop_sec short enough to round to zero samples raises a ValueError naming the argument instead of a bare ZeroDivisionError from inside the window math.

• Array-valued manifest columns enforce one kind per column: the ndarray branch of the column encoder checked only shape, so np.stack silently promoted mixed rows (int + float → float64, int + str → stringified '<U21'), violating the read-side contract that values keep the dtype they were written with; mixed kinds now raise naming the column, while widening within one kind (int32 + int64, <U3 + <U8) remains allowed and lossless. Trailing-NUL string/bytes leaves inside list rows are rejected like their scalar counterparts.

• get_stats()["total_batches"] counts batches again: the timestamp was minted once per item, so counting distinct timestamps returned the item count; write() now stamps the whole batch once.

Security

• Manifest-driven paths are confined to the dataset directory: a manifest travels with the data it describes, so a shared, mirrored or downloaded dataset hands the reader a file naming the paths it will open. TreeDataSource, AudioDataSource and the windowed-LUFS cache reader all concatenated a manifest-supplied filename onto a base directory with no validation, and pathlib discards the left operand when the right is absolute (Path("/data") / "/etc/passwd" is /etc/passwd) and does not normalize ... A tampered manifest.json therefore returned the contents of any file the training process could read as waveform, from where it flows into the model, the checkpoints and any pre-render output. All such joins now go through a guard that rejects absolute paths, .. segments, and anything resolving outside the dataset directory (including via a symlink). The rest of the manifest is validated too — num_samples, each leaf’s shape_per_sample and dtype (allow-listed to fixed-width numeric types), and every children key against the real AudioTree fields, which were previously splatted into the constructor as-is.

• Manifests are no longer unpickled: both readers loaded manifest.npz with allow_pickle=True. A manifest travels with the data it describes, so anyone who shares, mirrors or downloads a pre-rendered dataset hands the reader a file that gets unpickled — arbitrary code execution in every data worker. The flag was on only because string columns were stored as object arrays; the new audiotree._manifest stores them as fixed-width <U (bytes as |S) with an explicit __mask_<column> presence mask, so allow_pickle=False works. A side effect worth having: absence is now carried by the mask alone, so a stored -1, NaN or "" round-trips as itself where the old reader dropped empty tag cells. A pre-1.0 pickled manifest is refused by name rather than executed.

• AudioTree.from_manifest no longer opens arbitrary files: it joined manifest-supplied filenames with audio_dir / filename, and pathlib discards the left operand when the right is absolute and does not normalize .. — so a tampered manifest returned the contents of any readable file as waveform. The 1.0 traversal guard covered sources/ and missed this call site, because from_manifest was a second, divergent reader of the same file; it is now expressed on the shared parser and routed through _fs.safe_join. safe_join’s .. check also consulted only PurePosixPath, to which ..\..\etc\hosts is one opaque part; it checks both flavours now.

• Breaking — ExcerptConfig.search is resolved by name, not eval(): loudest_excerpt called eval() on a field deliberately typed str so argbind can bind it from YAML, making a config file arbitrary code execution in every data worker. It was not even usable as an extension point, since eval ran in audiotree.core’s namespace and only the two built-ins ever resolved. Names are now looked up in a registry ("uniform", "bias_early") with an importlib dotted-path fallback ("mypkg.offsets.my_search"), and resolution happens in __post_init__ so a typo raises when the config is built rather than minutes into training. The default is "uniform". The pre-1.0 spellings ("SaliencyParams.search_uniform", "SaliencyParams.search_bias_early") are not carried over — SaliencyParams itself became ExcerptConfig and search_function became search — so an existing config has to be updated to the short names.

[0.2.0] - 2025-02-17

• jit has been removed in most places. We encourage users to jit as late as possible.

• New class: AudioDataBalancedDataset, which is a grain Dataset, not a Data Source.

• AudioTree has a .latents property.

• New transform: NeuralLatentEncodeTransform.

• Class NeuralAudioCodecEncodeTransform has been adjusted. The arg is now encoder_fn and it takes an AudioTree instead of an audio data array.

• In an AudioTree’s metadata, the offset and duration will now be 1D arrays instead of 0D arrays.

• cpu has a kwarg has been removed in most places. You should think of AudioTrees as existing on CPU by default. If you pass them to a jitted function then they will be put on device.

• In, AudioDataSimpleSource and AudioDataBalancedSource, num_steps arg is now num_records. Also ._filepaths property is now .filepaths.

[0.1.0] - 2024-08-22

Breaking changes:

• SaliencyParams has moved from audiotree.datasources.SaliencyParams to audiotree.SaliencyParams

Updates:

The code has been tested with device parallel sharding. See the recent updates to DAC-JAX. SaliencyParams has a new search_function parameter. The two valid strings are SaliencyParams.search_uniform and SaliencyParams.search_early_bias. You can also plug in your own Callable function.

[0.0.5] - 2024-08-08

First release.