SpecuX

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

SpecuX provides FFTs and convolution, spectral transforms and reconstruction, audio features, filtering, and optional audio I/O through one functional API for NumPy, Torch, CuPy, and JAX. Results use the input array library. Resident arrays remain on their device.

Documentation: https://specux.com

import specux
import torch

x = torch.randn(8, 32768, device="cuda", requires_grad=True)
S = specux.stft(x, n_fft=1024, hop_length=256, output="power")
S.mean().backward()

assert S.device == x.device
assert x.grad is not None

Execution model

Array libraries, compute engines, graph bridges, and memory transports are separate concerns:

backend=None follows the input placement. An explicit backend= may stage a NumPy array through another engine. Resident Torch, CuPy, and JAX arrays are not redirected through the host. Move the array explicitly to change devices. XLA controls JAX placement.

ROCm is represented as a separate engine rather than treated as CUDA: the specux._rocm provider builds the same codegen kernels through HIP and hipRTC. It serves resident PyTorch-ROCm and CuPy-ROCm arrays (whose device type is still reported as cuda) by pointer, and stages NumPy through backend="rocm". PyTorch-ROCm is disambiguated from NVIDIA by torch.version.hip, so a HIP tensor is never routed to the CUDA engine. JAX-on-ROCm runs through the same XLA FFI bridge under the ROCM platform, so a resident JAX array stays on the GPU across jit, vmap, and gradients.

Features

• FFT and convolution: fft, ifft, rfft, and irfft at power-of-two, smooth, and prime lengths. Direct and FFT convolution support gradients for both operands.
• Transforms and reconstruction: stft, istft, cqt, vqt, icqt, and griffinlim, with complex, magnitude, power, and dB output modes.
• Features (specux.feature): mel spectrograms, MFCC, LFCC, chroma, RMS energy, onset strength, Tonnetz, spectral descriptors, deltas, inverse mel projection, and filterbank builders.
• Filtering (specux.filters): lfilter, filtfilt, differentiable biquad designs, preemphasis, and deemphasis. Signal and coefficient gradients are supported for lfilter.
• Audio I/O (specux.audio, optional): WAV, FLAC, MP3, OGG, and MP4 decoding and encoding; frame-accurate slicing; batch and streaming I/O; resampling; loudness, true-peak, and loudness-range metering; and metadata.
• Torch integration: resident eager kernels, torch.library custom ops, analytic adjoints, autocast policies, fake kernels, and torch.compile support. specux.transforms provides nn.Module wrappers.
• CuPy integration: resident CUDA kernels use the array's device pointer and current stream without host staging.
• JAX integration: typed XLA FFI on CPU, CUDA, and ROCm, including jax.jit, gradients, and batching. Elementwise dB operations remain fusible JAX expressions. JAX Metal is not supported.

Float16 uses float32 compute by default. Eligible GPU transforms may use native half compute with compute="half". Metal computes in float32 and does not support float64. JAX has no complex32, so half inputs widen to float32.

Install

pip install specux            # NumPy runtime; platform wheels include native engines
pip install specux[torch]     # Torch autograd and compile integration
pip install specux[jax]       # JAX runtime for the typed XLA bridge
pip install specux[cuda]      # Recommended CUDA runtime, currently CUDA 12

pip install specux[cuda11]    # Explicit CUDA 11 runtime
pip install specux[cuda12]    # Explicit CUDA 12 runtime
pip install specux[cuda13]    # Explicit CUDA 13 runtime
pip install specux[rocm]      # AMD HIP runtime and hipRTC (rocm-sdk wheels)

SpecuX uses one CUDA engine for CUDA 11, 12, and 13. The optional extras only install a matching NVRTC compiler and runtime headers. Use cuda for the recommended version or a numbered extra to select one explicitly. cuda remains an alias for cuda12 throughout the SpecuX 0.2 release line. The NVIDIA driver remains a system dependency.

No vendor library is linked: the CUDA driver and NVRTC, and the HIP runtime and hipRTC, are all resolved when an engine is first used. One wheel therefore carries the CPU, CUDA and ROCm engines, installs on a machine with none of them, and binds whichever are present - including both GPU vendors at once. The ROCm engine spans ROCm 6 and 7 from the same binary. AMD publishes the rocm-sdk wheels on its own index rather than PyPI, so specux[rocm] needs --index-url https://rocm.nightlies.amd.com/v2/gfx110X-dgpu/; a system ROCm install needs nothing extra.

specux.available("cuda")   # engine built in AND its runtime installed here
specux.devices()           # [Device(rocm:0 'AMD Radeon RX 7600S' gfx1102), ...]

Torch, CuPy, and JAX integrations are optional. Install the CuPy package that matches your CUDA environment. SpecuX does not choose or install a CuPy runtime. A base installation depends only on NumPy. Release wheels include the XLA bridge but do not install JAX.

Source builds

pip install -e .

An isolated build installs JAX in its temporary build environment to obtain the typed FFI headers. This does not add JAX to the installed runtime requirements. With --no-build-isolation, install JAX in the build environment before using SPECUX_JAX=1.

Useful build controls:

• SPECUX_JAX=0 omits the XLA bridge; SPECUX_JAX=1 requires it.
• SPECUX_CUDA=0 / SPECUX_ROCM=0 omit that GPU engine; =1 requires it. Both default to auto (build it when its SDK is found). The XLA bridge follows: it carries handlers for whichever GPU engines were built, and CUDA and ROCm can coexist in one build - they register under different XLA platforms - so a single wheel can serve both.
• SPECUX_CPU_ONLY=1 builds only the CPU engine.
• SPECUX_CUDA_HOME=/path/to/cuda selects a CUDA toolkit.
• SPECUX_AUDIO=0 omits audio I/O; SPECUX_AUDIO=1 requires FFmpeg headers.
• scripts/build.ps1 loads the MSVC environment and builds in place on Windows.

macOS source builds include the CPU and Metal engines. The Metal engine uses the vendored metal-cpp headers. Audio I/O requires FFmpeg development libraries, with TagLib used for tags and cover art. The DSP package remains usable when the audio extension is omitted.

Usage

NumPy and explicit staging

import numpy as np
import specux

x = np.random.default_rng(0).standard_normal((8, 32768)).astype(np.float32)

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

# NumPy uses CPU by default; an explicit engine stages the host array.
S_cuda = specux.stft(x, n_fft=1024, backend="cuda")

Audio I/O

import specux

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("session.flac", 30 * sr, sr=sr, mono=True):
    M = specux.feature.melspectrogram(block, sr=sr)

Torch compile and autocast

import specux
import torch

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

compiled = torch.compile(
    lambda v: specux.stft(v, n_fft=1024, hop_length=256, output="power"),
    fullgraph=True,
)

with torch.autocast("cuda", torch.float16):
    S = compiled(x)

CuPy resident CUDA

import cupy as cp
import specux

x = cp.random.standard_normal((8, 32768)).astype(cp.float32)
S = specux.stft(x, n_fft=1024, hop_length=256, output="power")

assert isinstance(S, cp.ndarray)
assert S.device.id == x.device.id

SpecuX launches on CuPy's current CUDA stream and keeps inputs, intermediates, and outputs on the device.

JAX JIT and gradients

import jax
import jax.numpy as jnp
import specux

x = jnp.ones((8, 32768), dtype=jnp.float32)

def loss(v):
    return specux.stft(v, n_fft=1024, hop_length=256, output="power").mean()

value, grad = jax.jit(jax.value_and_grad(loss))(x)

Enable jax_enable_x64 before creating arrays when using float64:

jax.config.update("jax_enable_x64", True)

Runtime controls

• specux.deterministic(True) selects reproducible overlap-add kernels and follows torch.use_deterministic_algorithms for Torch callers.
• specux.set_autotune(True) tunes new GPU configurations and persists launch wisdom.
• specux.autotune(...) explicitly tunes one operation and configuration.
• specux.clear_cache() releases clearable Python caches, native plans, and device tables.
• specux.clear_wisdom() removes persisted launch wisdom.

Design

SpecuX fuses framing, FFT, and output work when supported by the engine and transform shape. Multi-stage operations remain on the same device and stream, without intermediate copies through Python or host memory.

CUDA and Metal kernels share operation definitions through the accelerator code generator. The CPU engine implements the same operation families with AVX2 and NEON SIMD paths. Immutable windows, filterbanks, bases, and FFT data live in bounded native plans. SpecuX resolves those plans behind the functional API.

Torch integration uses a Python custom-op bridge, so the native engines do not depend on the libtorch ABI. CuPy passes resident buffers and the current stream directly to the CUDA engine. JAX integration uses a separate XLA FFI module, so JAX remains optional at runtime.

Native source layout and layering rules are documented in the native source guide.

Development

python scripts/lint.py
python -m pytest tests
python bench/bench_matrix.py

License

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