specux

Differentiable audio DSP for Python: fast, fused kernels on GPU and CPU.

Spectral transforms, an FFT family, and audio I/O, on numpy, torch, and cupy arrays alike.

Documentation: https://specux.com

import specux

y, sr = specux.audio.load("song.flac", sr=16000, mono=True)
M = specux.melspectrogram(y, sr=sr, n_fft=1024, n_mels=80)

import torch
x = torch.randn(8, 32768, device="cuda", requires_grad=True)
S = specux.stft(x, n_fft=1024, output="power")   # stays on the GPU
S.sum().backward()                                # native adjoint kernels

Arrays go in and come out in their own library: numpy in, numpy out; torch in, torch out (resident on its device, differentiable); cupy in, cupy out. The GPU kernels are generated in C++ and compiled at runtime (NVRTC for CUDA, MSL for Metal), so one build covers every size, precision, and output mode; the CPU backend carries the same transforms and adjoints, so training works without a GPU.

What's in the box

• Spectral transforms: stft / istft, melspectrogram, mfcc, lfcc, cqt / vqt / chroma, with output modes complex / magnitude / power / db and configured class twins (specux.STFT, specux.MFCC, ...). specux.transforms adds torch nn.Module wrappers with torchaudio-shaped defaults.
• FFT family: fft / ifft / rfft / irfft at any length (powers of two, smooth sizes, primes) and precision (float32, float64, and float16 storage with float32 compute), plus reusable plans (fft_plan, stft_plan with optional autotuning).
• Audio I/O (specux.audio): decode and encode WAV/FLAC/MP3/OGG/MP4, frame-accurate offset/duration, batch load_many/save_many, streaming readers and writers for multi-hour files, resampling, loudness / true-peak / loudness-range metering, and tags/cover metadata.
• torch integration: every entry point is a torch.library custom op, so torch.compile(fullgraph=True) traces without graph breaks and torch.autocast computes in float32. torch stays optional: without it, numpy arrays run on the CPU engine or the torch-free CUDA runtime.
• Runtime options: specux.deterministic(True) switches overlap-adds to bitwise-reproducible kernels (and follows torch.use_deterministic_algorithms); specux.benchmark(True) autotunes new configurations once and caches the result on disk.

Install

pip install specux              # CPU everywhere; CUDA/Metal where the wheel includes them
pip install specux[cuda12]      # + NVRTC and CUDA headers from NVIDIA's pip wheels
pip install specux[cuda13]      # the same for CUDA 13
pip install specux[torch]       # torch bundles its own CUDA, nothing extra needed

Or from source, in the environment you plan to use it in:

pip install -e .

The build is torch-free (no libtorch anywhere), so one build serves numpy, cupy, and whichever torch is installed at runtime. Python >= 3.10, numpy >= 1.22; torch optional (>= 2.4 for autograd and torch.compile,

= 2.7 for resident Metal).

• CUDA: builds when a toolkit is found (SPECUX_CUDA_HOME overrides); running needs an NVIDIA driver plus NVRTC from a toolkit, torch, or the nvidia-* pip wheels. Skipped with a notice otherwise.
• CPU: always builds. SPECUX_CPU_ONLY=1 forces a CPU-only build.
• macOS: builds the CPU and Metal extensions out of the box (metal-cpp is vendored).
• Audio I/O: needs FFmpeg dev libraries; on Windows, scripts/get_ffmpeg.ps1 fetches a self-contained LGPL build and scripts/get_taglib.ps1 adds tag/cover support. Without FFmpeg the audio module is skipped and everything else works.
• Windows: ./scripts/build.ps1 imports the MSVC environment and builds in place.

Usage

import numpy as np
import specux

x = np.random.randn(8, 32768).astype(np.float32)

# functional, any array library
S = specux.stft(x, n_fft=1024, hop_length=256, output="power")
y = specux.istft(specux.stft(x, 1024), 1024, length=x.shape[-1])
C = specux.mfcc(x, sr=16000, n_mfcc=20)

# configured twins: construct once, call with (..., time)
t = specux.MelSpectrogram(sr=44100, n_fft=1024, n_mels=80)
M = t(x)

# pick a backend explicitly (default follows the input)
S = specux.stft(x, n_fft=1024, backend="cpu")

import torch

xt = torch.randn(8, 32768, device="cuda")

# torch.compile and autocast
f = torch.compile(lambda v: specux.stft(v, 1024, output="power"), fullgraph=True)
with torch.autocast("cuda", torch.float16):
    S = specux.stft(xt, 1024)          # computes in float32

# reusable plan; tuning knobs live only here
plan = specux.stft_plan(n_fft=1024, hop_length=256, output="power", device="cuda")
plan = specux.autotune(xt, plan)       # optional, cached on disk
S = plan(xt)

# audio: read, meter, transform, write
y, sr = specux.audio.load("take.wav", sr=16000, mono=True)
lufs = specux.audio.loudness(y, sr)
y = specux.audio.normalize(y, mode="lufs", target_db=-14.0, sr=sr)
specux.audio.save("out.flac", y, sr)

for block in specux.audio.blocks("4hours.flac", 30 * sr, sr=sr, mono=True):
    M = specux.melspectrogram(block, sr=sr)     # constant memory

Design

A transform is one fused kernel. The C++ codegen (src/codegen/) assembles each kernel from three parts:

• a prologue that frames, reflect-pads, and windows the signal,
• an FFT core picked by size and precision,
• an epilogue that finishes the op in the same pass: the Hermitian recombine, then the output mode (complex / magnitude / power / db), a filterbank reduce for mel, a wavelet-basis reduce for the CQT family, or the spectrum product for convolution.

The spectrum never round-trips through memory between those stages: a dB mel spectrogram is one kernel, and features like MFCC or chroma are short chains of them. Every backward pass is the analytic adjoint of the same chain, generated the same way. Op, direction, and mode select the prologue and epilogue, precision is a type parameter, and size is a plan, so adding a size or mode never adds a hand-written kernel. Lengths past a GPU block's shared memory decompose into a two-kernel four-step pipeline, and large prime factors take a chirp-z (Bluestein) route through the same machinery.

The same sources emit the CUDA and Metal dialects, and the CPU engine (src/cpu/) implements the same transforms and adjoints. src/README.md describes the native layout and layering rules.

Tests cover correctness, autograd, torch.compile, autocast, and audio; every numerical tolerance lives in tests/_tol.py with its derivation.

python -m pytest tests        # correctness, autograd, compile, autocast, audio
python bench/bench_matrix.py  # timing matrix on your own hardware

License

MIT. Vendored third-party components and their licenses are listed in NOTICE.