mlx-sci 0.2.0

The missing scipy toolkit for Apple Silicon — GPU-accelerated special functions, linear algebra, signal processing, and quantum information via MLX

mlx-sci

The missing scipy + quantum-information toolkit for Apple Silicon. GPU-accelerated via MLX.

mlx-sci is a meta-package that bundles a curated set of focused sub-packages — Airy, Bessel, Gamma, Hypergeometric, Wigner, matrix exponential, STFT, Fisher information, quantum relative entropy, and a statevector circuit simulator — under one consistent namespace (mlx_sci.special, mlx_sci.linalg, mlx_sci.signal, mlx_sci.quantum). Each sub-package is independently installable, so you can either pull the whole stack or pick à la carte.

Everything runs on the Apple GPU through MLX. There is no CUDA path, no host->device shuffling, and no framework dependency beyond MLX + NumPy.

Install

pip install mlx-sci         # bundles all sub-packages

# or pick à la carte
pip install mlx-airy
pip install mlx-bessel
pip install mlx-expm
pip install mlx-fisher
pip install mlx-gamma
pip install mlx-hyp2f1
pip install mlx-qre
pip install mlx-stft
pip install mlx-wigner
pip install mlx-quantum-sim

Python >= 3.10, Apple Silicon (M1/M2/M3/M4), MLX >= 0.30.

Quick Start

import mlx.core as mx
from mlx_sci import special, linalg, signal, quantum

# Airy functions over a real grid
x = mx.linspace(-15.0, 15.0, 10_000)
Ai, Ai_prime, Bi, Bi_prime = special.airy(x)

# Gamma / lgamma / digamma -- vectorised on GPU
y = special.gamma(mx.linspace(0.1, 30.0, 1_000_000))

# Gauss hypergeometric, auto-routed to a fused Metal kernel
a = mx.array(0.5); b = mx.array(1.0); c = mx.array(1.5)
z = mx.linspace(0.01, 0.95, 1_000_000)
F = special.hyp2f1(a, b, c, z)

# Matrix exponential (Pade-13 + scaling-squaring, on GPU)
H = mx.random.normal((1024, 1024))
U = linalg.expm(-1j * H * 0.1)

# STFT of an audio buffer
audio = mx.random.normal((480_000,))             # 30 s @ 16 kHz
spec = signal.stft(audio, n_fft=1024, hop_length=256)

# Wigner 3j -- one million coupling coefficients in a single dispatch
import mlx.core as mx
j1 = mx.ones(1_000_000); j2 = mx.ones(1_000_000); j3 = 2 * mx.ones(1_000_000)
m = mx.zeros(1_000_000)
w3j = special.wigner_3j(j1, j2, j3, m, m, m)

# Quantum relative entropy: exact eigh path
rho   = quantum.random_density_matrix(256)
sigma = quantum.random_density_matrix(256)
D = quantum.quantum_relative_entropy(rho, sigma)

# Quantum relative entropy: Stochastic Lanczos for large N
# O(k * N^2), beats the eigh path by 1-2 orders at N >= 1000.
rho_big   = quantum.random_density_matrix(2000)
sigma_big = quantum.random_density_matrix(2000)
D_slq = quantum.quantum_relative_entropy_lanczos(
    rho_big, sigma_big, k=25, m=20
)

# Statevector circuit simulator
sim = quantum.MLXQuantumSimulator(3)
sim.h(0); sim.cx(0, 1); sim.cx(0, 2)
probs = sim.measure_probs()

Modules at a glance

Module	Source sub-package	Scope	One-liner
mlx_sci.special.airy	mlx-airy	Airy Ai/Bi + derivatives	All four outputs in one call.
mlx_sci.special.gamma	mlx-gamma	gamma, lgamma, digamma, beta	Vectorised on GPU.
mlx_sci.special.hyp2f1	mlx-hyp2f1	Gauss 2F1, 1F1, 0F1	Fused metal_kernel collapses ~200 MLX ops into 1 dispatch.
mlx_sci.special.BesselTable	mlx-bessel	Spherical Bessel j_l(x), j_l'(x)	Build once, evaluate over arbitrary x grids.
mlx_sci.special.wigner_*	mlx-wigner	Wigner 3j / 6j / 9j, Clebsch-Gordan	Racah formula on GPU; millions per call.
mlx_sci.linalg.expm	mlx-expm	matrix expm / logm / sqrtm / Frechet	Pade-13 + scaling-squaring on GPU.
mlx_sci.signal.stft	mlx-stft	STFT / ISTFT, mel filterbank, windows	Whisper-style mel spectrograms inclusive.
mlx_sci.quantum.qre	mlx-qre	D(rho || sigma), von Neumann entropy	Exact eigh + Stochastic Lanczos quadrature.
mlx_sci.quantum.fisher	mlx-fisher	Fisher information matrix, natural-grad	Cosmology-scale J^T W J on GPU.
mlx_sci.quantum.sim	mlx-quantum-sim	Statevector simulator, batched + noisy	Ideal + WILLOW / HERON / T9 noise profiles.

Performance

All numbers below were measured on an Apple M1 Max, MLX 0.30-0.31, NumPy 2.x, SciPy 1.16. SciPy / NumPy reference is float64 on the CPU (Accelerate / LAPACK); MLX paths are float32 on the Apple GPU.

Module	What it accelerates	Headline speedup (M1 Max)	Break-even
mlx_sci.special.BesselTable	Spherical Bessel j_l (eval-only)	579x @ N_ell=525, N_x=10k	N_x >= 5k (or table re-used)
mlx_sci.special.airy	Airy Ai/Bi + derivatives	6.7x @ N=1M	N >= ~30-50k
mlx_sci.special.gamma	gamma / lgamma / digamma	gamma 10.4x, lgamma 6.4x @ N=1M	N >= ~50k
mlx_sci.special.hyp2f1	Gauss 2F1 (fused Metal kernel)	7.4x @ N=1M	N >= ~100k
mlx_sci.linalg.expm	Matrix exponential (Pade-13)	2.08x real @ n=1024, 1.96x complex @ n=256	n >= 1024 real / n >= 256 complex
mlx_sci.special.wigner_3j	Wigner 3j/6j/9j (Racah)	2000x @ batch=1k vs sympy	batch >= ~1k
mlx_sci.signal.stft	STFT / mel spectrogram	10.0x @ 30 s audio (16 kHz)	duration >= ~5 s
mlx_sci.quantum.fisher	Fisher J^T W J / large matmul	39x @ 32k x 512	matrix size dependent
mlx_sci.quantum.qre (eigh)	Quantum relative entropy	1.93x @ N=1000	N >= ~500
mlx_sci.quantum.qre (Lanczos)	QRE via Stochastic Lanczos	657x @ N=2000 vs NumPy exact, 84x @ N=1000	N >= 1000 (exact times out at N=2000)
mlx_sci.quantum.sim	Statevector simulator	see sub-package README	qubit-count dependent

See each sub-package's benchmark_results.md for the full curve, break-even points, and accuracy tables.

Why MLX, not CUDA / NumPy?

• Apple Silicon native. No CUDA, no ROCm, no CPU offload dance — the Metal GPU on your laptop is the same one this stack is benchmarked on. No external dependencies beyond MLX + NumPy.
• Lazy evaluation. MLX builds a deferred graph; the entire computation is materialised once at mx.eval time. We exploit this in mlx-qre's Stochastic Lanczos hot path so the whole k-step recurrence becomes a single GPU command-buffer.
• mx.compile fusion. Fusable element-wise + reduction subgraphs are JIT-compiled into single kernels.
• mx.fast.metal_kernel for hot inner loops. When auto-fusion cannot collapse a 200-op Taylor series into one dispatch, we drop in a hand-written Metal kernel — the 7.4x hyp2f1 speedup over SciPy comes from exactly this trick (the underlying *_metal symbols are internal; users keep calling hyp2f1/hyp1f1/hyp0f1 and the routing is transparent).
• Honest about break-even. Every sub-package documents the N below which SciPy on CPU still wins. We do not claim wins where we lose; we publish the cross-over point and recommend the right tool for the size.

API stability

mlx-sci is at v0.2.0. The 0.x series is still iterating — minor versions may add re-exports and refactor module layout. We commit to keeping the headline functions (airy, gamma, hyp2f1, expm, stft, quantum_relative_entropy, MLXQuantumSimulator) source- compatible across the 0.x line. A v1.0 release will lock the public surface.

Sigma = 2 ln Q ecosystem

mlx-qre and mlx-quantum-sim are the numerical backbone for a small constellation of physics projects organised around the identity Sigma = 2 ln Q:

• anatropic — entropy production / second-law violation tests on near-term quantum hardware.
• petz-recovery-unification — Petz recovery map experiments and fidelity bounds.
• tau-chrono — time-symmetry / Khronon-foliation experiments on the IQM Tuna-9 / Tuna-17 backends.

If you only care about the numerics, you can ignore that side entirely; the sub-packages are framework-agnostic.

Citation

If mlx-sci saves you time, please cite either the meta-package or the specific sub-package(s) you actually used:

@software{mlx_sci,
  author  = {Huang, Sheng-Kai},
  title   = {mlx-sci: GPU-accelerated SciPy + quantum-information for
             Apple Silicon},
  year    = {2026},
  url     = {https://github.com/akaiHuang/mlx-sci},
  version = {0.2.0},
}

Each sub-package (mlx-airy, mlx-bessel, mlx-expm, mlx-fisher, mlx-gamma, mlx-hyp2f1, mlx-qre, mlx-stft, mlx-wigner, mlx-quantum-sim) ships its own short BibTeX entry in its README; cite the one(s) you actually exercised, not the umbrella, where the distinction matters.

License

MIT.