STFT

The Short-Time Fourier Transform slices a signal into hopped, windowed frames and takes the real FFT of each: the spectrogram every other transform on this site builds on. It runs as one fused kernel on the GPU backends (window, real FFT, and the output policy never leave the chip), and any n_fft works: powers of two, 2/3/5-smooth sizes, primes, and everything else each get a specialized kernel. The inverse is ISTFT.

specux.stft(x, plan=None, n_fft=2048, hop_length=None, win_length=None,
            window="hann", center=True, output="complex", eps=1e-10,
            backend="auto")

numpy

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")
S.shape        # (8, 513, 129) = (..., n_fft // 2 + 1, n_frames)

torch

import torch

import specux

xt = torch.randn(8, 32768, device="cuda", requires_grad=True)
S = specux.stft(xt, n_fft=1024, hop_length=256, output="power")
S.sum().backward()             # differentiable, resident on the GPU
xt.grad.shape                  # torch.Size([8, 32768])

cupy

import cupy

import specux

xc = cupy.random.standard_normal((8, 32768), dtype=cupy.float32)
S = specux.stft(xc, n_fft=1024, hop_length=256, output="power")
S.shape        # (8, 513, 129), on the GPU, on cupy's own stream, no torch

Parameters

• x: the signal(s), shaped (..., time). numpy, torch, or cupy; the result comes back in the same library on the same device.
• plan: a prebuilt specux.stft_plan(...) handle for hot loops with fixed parameters; None plans automatically. See Plans & autotuning.
• n_fft: transform size, any integer.
• hop_length: frame advance in samples; defaults to n_fft // 4.
• window: "hann" (default), "hamming", "blackman", "blackmanharris", "nuttall", "bartlett", "boxcar", or your own length-n_fft array. The leakage trade-offs between them are classic [harris1978].
• center: reflect-pad by n_fft // 2 so frame t is centered on t * hop_length (default True).
• output: "complex" (default), "magnitude", "power", or "db". Fused into the kernel: for the real modes the complex spectrum never touches global memory.
• backend: "auto" (follows the input’s device), "cuda", "cpu", or "metal".

win_length must equal n_fft (or stay None) in v1, and eps floors the "db" mode.

The configured form

specux.STFT holds the parameters and broadcasts over any leading shape; to_dict() round-trips the configuration:

t = specux.STFT(n_fft=1024, hop_length=256, output="power")
t(np.random.randn(2, 3, 32768).astype(np.float32)).shape   # (2, 3, 513, 129)
assert specux.STFT(**t.to_dict()) == t

For training pipelines, specux.transforms.Spectrogram is the same transform as a torch.nn.Module (window registered as a buffer, nn.Sequential-ready). For hot loops with fixed parameters, see Plans & autotuning.

References

• [allen1977]: the analysis/synthesis framework this implementation follows.
• [harris1978]: window functions and their leakage trade-offs.