mlx-signal

GPU-accelerated signal processing for Apple Silicon, with familiar scipy.signal APIs.

mlx-signal implements a practical subset of scipy.signal with MLX and custom Metal kernels. It accepts NumPy or MLX arrays and returns MLX arrays in Apple's unified memory.

import numpy as np
import mlx_signal_processing as sig

x = np.random.randn(64, 1 << 20).astype(np.float32)
frequencies, power = sig.welch(x, fs=48_000, nperseg=1024)
power_np = np.array(power)

MLX pipelines can pass results straight into a model without a host-device transfer. NumPy users get the same API and can convert the result with a normal in-memory copy. Small jobs automatically stay on SciPy when a GPU launch would cost more than it saves.

Install

Requires Apple Silicon, macOS 14 or newer, and Python 3.10 or newer:

python -m pip install mlx-signal-processing

The distribution is named mlx-signal-processing; import it in Python as mlx_signal_processing.

To work on mlx-signal from a checkout:

git clone https://github.com/tabulai/mlx-signal
cd mlx-signal
python -m pip install -e .   # or: uv pip install -e .
python -m pytest -q         # optional: run the SciPy parity suite

Measured performance

Measured on an Apple M4 Max (macOS 26.2) with MLX 0.32.2, SciPy 1.18.1, and float32 data. Every result passed a SciPy correctness check before timing. Values are medians of 9 runs after 3 warmups:

• e2e: NumPy input and output, for drop-in use
• device: MLX input and output, for an on-device pipeline

See the full report, or reproduce it with python bench/bench.py --warmup 3 --repeat 9.

¹ MLX 0.32's Metal FFT fails at some lengths above 2^20. mlx-signal handles them with its own four-step (Bailey) GPU decomposition. Power-of-two lengths from 2^21 through 2^26 use three fused Metal passes: about 2x faster than the composed path at 2^23 and 5x at 2^26, with better accuracy than MLX's large FFT. Other factorable lengths use safe-size sub-FFTs; only lengths without a safe factorization, such as large primes, use the CPU stream. See Known limitations.

² The expensive prominence search in find_peaks runs on the GPU and matches SciPy bit for bit. Index bookkeeping stays on the host, limiting the overall speedup to about 2.5x.

Comparison with other libraries

This end-to-end comparison uses the same machine and NumPy input/output. Shapes and conventions are aligned, and every result is checked against SciPy. The fastest result in each row is bold. See the full report, or reproduce it with uv sync --extra bench && uv run python bench/bench_cross.py.

¹ Torchaudio's default anti-aliasing filter is much shorter (lowpass_filter_width=6, versus 3201 taps here). mlx-signal keeps SciPy's default filter and matching output, yet still wins end to end. With arrays already on the GPU, it is 3.4x faster (0.99 versus 3.34 ms).

² Torchaudio offers FFT convolution, but lfilter does not select it automatically. The closest causal FIR operation is conv1d (O(n·k)), used in the table. Torchaudio's general lfilter takes 3.4 s on CPU and 22.4 s on MPS for this case—more than 1600x slower than mlx-signal.

What the comparison shows:

• In this group, mlx-signal is the only GPU implementation of welch/csd/coherence. JAX runs those APIs on CPU, and torchaudio has no PSD estimator. The fused two-signal kernel computes coherence in one pass, compared with five in SciPy.
• upfirdn and find_peaks have no comparable implementation in the other libraries tested.
• torch.stft on MPS is close end to end (4.35 versus 4.28 ms) because data transfer dominates and the NumPy-to-GPU copy time varies. On-device, mlx-signal is 2.7x faster (0.92 versus 2.47 ms), without pulling in the 2 GB torch dependency.

The spectral hot path uses fused Metal kernels. STFT/Welch read each segment once without materializing a frames array; ISTFT combines an inverse transform with gather-based overlap-add. Shapes that do not fit the fused path use zero-copy framing and compiled MLX operations.

What's implemented (v0.1)

IIR filtering

lfilter, filtfilt, sosfilt, sosfiltfilt, and the default IIR path in decimate run on the GPU, including single-channel inputs. SOS and transfer-function filters support native zi/zf, so state can carry across chunks. Transfer-function filters are supported through order 16.

Long signals use a block-parallel scan; shorter jobs use a sequential kernel when worthwhile, then fall back to SciPy. In the table above, transfer-function filtering reaches 128x at 256 channels and 27x for a single channel on-device. The sequential path is bit-identical to matching SciPy float32 builds. The parallel path typically matches to about 1e-6, with roughly 1e-5 worst-case error for resonant filters near its routing threshold.

High-order filters with clustered poles stay on the sequential path. Complex coefficients and transfer-function orders above 16 fall back to SciPy with a FallbackWarning; use SOS form for those filters.

Dispatch: when the GPU is used

GPU launches do not pay off for small inputs, so mlx-signal can choose the backend for each call:

• dispatch="auto" (default) uses MLX above gpu_min_size (2^15 work elements by default) and SciPy below it.
• dispatch="mlx" always uses MLX and raises NotImplementedError when no MLX path exists.
• dispatch="scipy" uses SciPy's numerical kernels with canonical float32/complex64 inputs. It still returns MLX arrays and does not restore float64 SOS arithmetic.

All three modes return the same MLX types and dtypes.

sig.set_config(dispatch="mlx")                  # global
with sig.config_context(gpu_min_size=1 << 18):  # scoped
    ...

Capability fallbacks issue a warning. Examples include complex IIR coefficients, transfer-function orders above 16, callable detrending, signals too short for a requested boundary extension, and exceptional non-finite filter or padding cases. Size-based routing is silent.

Dtype policy

• Computation uses float32/complex64 because Metal does not support float64.
• Explicit float64/complex128 signal and state arrays downcast with a one-time DowncastWarning. Use set_config(float64="strict") to raise instead, or warn_on_downcast=False to disable the warning. Extended-precision SOS arrays follow the same rule.
• The parity suite compares the float32 pipeline with float64 SciPy references at rtol=1e-4 and a peak-relative atol≈1e-5.
• SciPy's filter-design functions normally return small float64 SOS arrays. These convert quietly unless strict mode is enabled. Filtering stops with an error if conversion makes a design unstable or erases a section numerator.
• With SciPy older than 1.15, unsafe automatic SOS scans stay on SciPy because its historical float32 recurrence order differs from the Metal kernel. Transfer-function lfilter probes the installed SciPy build and takes the same conservative route when its compiler-dependent rounding differs.
• Apple GPUs flush float32 denormals below about 1.2e-38 to zero; CPU SciPy keeps them. IIR bit-identity therefore applies to normal-range data.

Known limitations

• MLX 0.32 has an upstream Metal FFT issue above 2^20. Lengths in (2^19, 2^21], except 2^20, crash with Unable to load function four_step_mem_…; other lengths above 2^20 can return incorrect values with relative error around 1.0. mlx-signal uses its own four-step decomposition for 1-D transforms and blocked overlap-add for long×short convolution. Unsplittable lengths such as large primes, and N-D FFT paths, use MLX's CPU stream. The workaround can be removed when MLX fixes the affected range.
• Windows/filter design (get_window, firwin*) and find_peaks' index refinement run on the host. The prominence search itself runs on the GPU for float32 inputs.
• lfilter/sosfilt take and return zi/zf state natively on the GPU for IIR filters. FIR lfilter with zi uses SciPy's CPU convolution path and warns.
• upfirdn handles every SciPy extension mode on-device. It falls back for signals shorter than the required extension (roughly the tap count divided by up, adjusted for downsampling phase), and for non-finite taps with exceptional NaN/Inf edge semantics.
• Streaming APIs, ShortTimeFFT, and CWT are not yet implemented.

Roadmap

• CWT (removed from SciPy in 1.15)
• Modern ShortTimeFFT class
• float16 mode
• Real-time streaming API
• Torchaudio benchmark coverage

Examples

• examples/fm_demod.py — an SDR FM demodulation chain (channel filter → polyphase decimate → discriminator → de-emphasis → audio resample), 16.5x end-to-end versus SciPy on an M4 Max, with 0.999 correlation to the true message.
• examples/eeg_bandpower.py — 64-channel × 10-minute EEG alpha-band power via one batched welch, 6x versus SciPy including result readback.

Development

uv venv && uv pip install -e ".[dev]"
python -m pytest -q          # parity tests against SciPy
ruff check src tests
python bench/bench.py        # GPU benchmark

CI runs lint and the full test suite on GitHub's arm64 macOS runners. Tests marked gpu require a Metal device; the rest also exercise MLX's CPU backend.

Acknowledgments

• SciPy provides the API contract and numerical reference. Edge cases are matched against scipy.signal (BSD-3-Clause) in the parity suite.
• MLX provides the lazy, unified-memory array framework.