specux

Differentiable STFT, iSTFT, and mel spectrograms: one array-agnostic API over fused CUDA and Metal kernels generated at runtime, with a fast CPU backend everywhere else.

Quickstart GitHub

mel spectrogram of a rising three-harmonic chirp, rendered by specux

pip install specux              # CPU everywhere; Metal on macOS, CUDA with a driver + NVRTC
pip install specux[cuda12]      # + NVRTC & CUDA headers from NVIDIA's pip wheels
pip install specux[cuda13]      # the same for CUDA 13 builds
pip install specux[torch]       # torch bundles its own CUDA, no extra needed

import numpy as np

import specux

x = np.random.randn(8, 32768).astype(np.float32)
S = specux.stft(x, n_fft=1024, hop_length=256, output="power")   # (8, 513, 129)
M = specux.melspectrogram(x, sr=16000, n_fft=1024, n_mels=80)    # (8, 80, 129)

torch tensors stay on their device and differentiate:

import torch

xt = torch.randn(8, 32768, device="cuda", requires_grad=True)    # or "mps"
specux.melspectrogram(xt, sr=16000, n_fft=1024, n_mels=80).sum().backward()

Round trips are exact to float precision:

Z = specux.stft(x, n_fft=1024, hop_length=256)
y = specux.istft(Z, n_fft=1024, hop_length=256, length=x.shape[-1])
np.abs(y - x).max()                                              # ~1e-5

Any n_fft, f32 / f64 / f16 storage, numpy / torch / cupy in and out. The Quickstart covers install and the rest.