auvux-dsp 0.1.0.dev4

Fast differentiable audio transforms (STFT, mel, MFCC, CQT, chroma) on CPU and GPU

auvux-dsp

Fast differentiable audio transforms (STFT, iSTFT, mel, MFCC, CQT/VQT, chroma) plus differentiable scalar/temporal features (spectral moments & flatness, RMS, spectral flux, onset strength, tempogram) with native CPU backends (vDSP on macOS, PFFFT elsewhere) and GPU support (Metal, CUDA). Forward and backward passes run in native kernels; PyTorch autograd plugs in when you pass a torch tensor, and torch is never required otherwise.

import auvux.dsp as dsp

mel = dsp.MelSpectrogram(sr=44100, n_fft=2048, hop_length=512, n_mels=128)
S = mel(y)                       # numpy in -> numpy out
S = mel(y, backend="gpu")        # Metal / CUDA kernels

y = torch.tensor(clip, requires_grad=True)
loss = mel(y, output="db").sum() # native forward
loss.backward()                  # native adjoint kernel, no torch recompute

pip install auvux-dsp --pre    # preview release; drop --pre once 0.1.0 is out

Performance

Benchmarks against librosa and torchaudio (forward passes and full training steps, CPU / staged GPU / GPU-resident) live in benchmarks/benchmark.py:

python benchmarks/benchmark.py

Results table to be published with the first release.

Complex STFT layout

Complex spectra are returned as (..., bins, frames) backed by frame-major memory — each frame's spectrum contiguous, the freq axis strided. This is the same physical layout both references use (librosa allocates its stft output order='F'; torch.stft returns a transposed view over frame-major memory), so values and bytes match librosa, and the GPU-resident path returns a tensor with the exact strides torch.stft produces. It is also what makes the STFT fast: no backend ever materializes the bins-major transpose. istft accepts both this layout (zero-copy) and compact C-order arrays. Float outputs (power/db/mel/...) are ordinary C-contiguous arrays.

Dtypes

Outputs follow the input dtype, torch-style:

input	CPU	GPU (Metal / CUDA)
float32 / ints (promote)	native float32	native float32
float64 / complex128	native double engines (validation grade: torch.autograd.gradcheck passes on every adjoint; librosa-double parity ~1e-10)	backend="auto" routes to CPU; explicit GPU backend raises
float16 (torch, GPU-resident)	computed in float32, cast back	native half kernels (Metal; CUDA planned): f16 in → f16/complex32 out, FFT arithmetic in half with float accumulators — torch.stft-CUDA-half numerics (~1e-3) at half the bandwidth
float16 (numpy / staged) / bfloat16	computed in float32, real outputs cast back to the input dtype	same (upcast on device)

torch CPU rejects half-precision FFTs outright and torch cannot create float64 MPS tensors, so this is a superset of torch.stft's dtype matrix. Complex outputs for staged half inputs stay complex64; the native Metal path returns torch.complex32 exactly like torch.stft on CUDA half (normalize signals to ~±1 — half overflows at 65504, like every fp16 pipeline).

FFT sizes

Any even n_fft works, torch-style, picked per size at plan time: powers of two run radix-2 (vDSP on macOS, PFFFT elsewhere), other 2^a·3^b·5^c sizes (96, 384, 960, …) run PFFFT mixed-radix at full speed, and everything else (400, 1000, primes·2, …) falls back to a Bluestein chirp-z on the pow2 engines (~5× a direct FFT — the same penalty torch pays inside pocketfft/cuFFT for such sizes). Whisper's n_fft=400 / hop=160 / 80 mels works out of the box. On GPU, pow2 sizes use the Stockham kernels and any other even size in [64, 4096] uses direct-DFT kernels; gradients flow through every path.

Status: under construction.

• CPU (vDSP/PFFFT): STFT, ISTFT, MelSpectrogram, MFCC, CQT, VQT, Chroma — forward and native backward, librosa-parity tested, torch autograd built in.
• Differentiable features: SpectralMoments (centroid, bandwidth, skewness, kurtosis, flatness, slope, decrease, crest, entropy), RMS, SpectralFlux, OnsetStrength, Tempogram — native forward + analytic adjoint, librosa-parity / torch-grad tested.
• Composed descriptors (built on the transforms above, so they inherit the GPU path and stay differentiable): SpectralContrast, BandEnergy, Tonnetz, ChromaCENS, FourierTempogram. Plus delta (Savitzky-Golay derivative) and pool post-processing over any feature matrix.
• Metal: every native transform on GPU (n_fft <= 4096), forward + backward, parity-tested against the CPU path. torch MPS tensors stay on the GPU end to end (DLPack), and backend="auto" routes them there — no flags needed.
• CUDA: kernel-for-kernel twin of the Metal backend including the resident paths; parity-tested on NVIDIA hardware (RTX 4090, CUDA 12.9), with pinned double-buffered staging for the numpy-in/numpy-out GPU paths.
• Pending: iCQT.

Development

pip install scikit-build-core pybind11 numpy pytest
./scripts/dev-build.sh
pytest
pytest -m soak          # memory-plateau soak (deselected by default)
./scripts/sanitize.sh   # ASan+UBSan+LSan suite; Linux, or via docker (see header)

Note for packagers: python/auvux/ is a PEP 420 namespace package — it must never contain an __init__.py, or it will shadow sibling auvux-* distributions.