Backends

Three backends, one API: CUDA (kernels generated and compiled at runtime via NVRTC), Metal (same codegen, compiled as MSL on Apple GPUs), and CPU. backend="auto" (the default everywhere) routes by the input:

• a torch tensor runs where it lives: CUDA tensor → CUDA kernels, mps tensor → Metal kernels, CPU tensor → CPU backend. The result stays on that device.
• a cupy array runs on CUDA by raw device pointer, on cupy’s own current stream; torch is not involved.
• a numpy array runs on CUDA when a CUDA device is present (with torch, staged through pinned memory; without torch, the extension stages it itself), otherwise on the CPU backend. On macOS, pass backend="metal" to run numpy arrays on the Apple GPU.

Passing backend= explicitly is strict: backend="cpu" with a GPU tensor raises instead of silently copying through the host.

Support matrix

	CUDA	Metal	CPU
input libraries	numpy, torch, cupy	numpy, torch	numpy, torch, cupy
precisions	f32, f64, f16	f32, f16	f32, f64
gradients (torch)	✓	✓	✓
torch.compile	✓	✓	✓
plans & autotune	✓	✗	✗
deterministic mode	✓	✓	✓
requires	NVIDIA driver + NVRTC	macOS	numpy

Notes on the fine print:

• f16 is a storage format: on both GPU backends, half inputs load and store half but compute in float32 registers.
• Any n_fft, any length: transforms that fit one GPU block run fused; larger ones (any n_fft, past a hundred thousand points and beyond) switch to a multi-kernel factor pipeline on the device, so no size falls back to the CPU on CUDA.
• CUDA works without torch: cupy arrays pass device pointers and their current stream straight to the kernels, and numpy inputs stage host to device inside the extension itself (the native route). With torch installed, numpy gets the faster transport (pinned-memory staging, torch’s allocator and streams) and the plan cache rides on it.
• Metal is torch-free for numpy: backend="metal" with a numpy array stages through host memory inside the extension itself.
• Resident Metal has two transports: from torch 2.9, mps tensors cross zero-copy via DLPack (the primary path; it keeps specux’s merged command buffers). On torch 2.7/2.8, which lack MPS DLPack export, the same generated MSL dispatches through torch.mps.compile_shader on torch’s own stream: every transform and feature, all output modes, pow2 and non-pow2, forward and backward, so training runs resident on those releases too.
• cupy has no autograd of its own, so cupy inputs are forward-only by nature; everything else (all sizes, all output modes, non-contiguous inputs) matches the numpy path bit for bit on the same backend.
• NVRTC comes from a CUDA toolkit, from torch’s bundled libraries, or from the specux[cuda12] / specux[cuda13] pip extras, whichever is found at runtime.

Gradients

torch inputs with requires_grad are differentiable through every transform and feature on every backend: stft, istft, melspectrogram, mfcc, lfcc, cqt, vqt, chroma, fftconvolve, and the FFT family. The backward pass is the analytic adjoint: a fused kernel chain of the same family as the forward, not autograd tracing through a decomposition, so training steps stay at kernel-launch cost:

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

The gradients are verified against torch.autograd.gradcheck in float64, and CPU↔CUDA gradient parity is part of the test suite. Inside torch.compile, every call on every backend routes through dispatcher-registered ops so the graph traces whole, forward and backward, with zero graph breaks (fullgraph=True works); inside torch.autocast the transforms compute in float32 regardless of the surrounding dtype (the AMP-safe policy for spectral ops).

Determinism in training: the backward of stft scatters overlapping frames. specux.deterministic(True) (or torch’s use_deterministic_algorithms) makes it bitwise-reproducible; see Runtime options.