mlx-wigner 0.1.0

GPU-accelerated Wigner 3j/6j/9j symbols via Apple MLX — the world's first GPU implementation

mlx-wigner

GPU-accelerated Wigner 3j/6j/9j symbols for Apple Silicon via MLX. Fully vectorized -- evaluate millions of symbols in a single call.

Why

Wigner 3j/6j/9j symbols appear everywhere in quantum mechanics, atomic physics, and nuclear structure calculations. Every existing implementation (sympy, wigxjpf, SHTOOLS) is CPU-only and computes symbols one at a time. For modern applications that require millions of coupling coefficients (e.g., spherical tensor networks, many-body perturbation theory, angular power spectra), serial CPU evaluation is the bottleneck.

mlx-wigner solves this by:

• GPU-vectorized Racah formula — all symbols in a batch are computed simultaneously
• Log-gamma numerics — stable for angular momenta up to j ~ 100
• Zero CPU loops — the entire computation graph runs on GPU

Installation

pip install mlx-wigner

Or from source:

git clone https://github.com/akaiHuang/mlx-wigner.git
cd mlx-wigner
pip install -e ".[dev]"

Requires Python >= 3.10 and Apple Silicon (M1/M2/M3/M4).

Quick start

from mlx_wigner import wigner_3j, wigner_6j, wigner_9j, clebsch_gordan

# Single symbol
w = wigner_3j(1, 1, 2, 0, 0, 0)

# Clebsch-Gordan coefficient <j1 m1; j2 m2 | J M>
cg = clebsch_gordan(1, 0, 1, 0, 2, 0)

# Batch: 100,000 symbols at once on GPU
import mlx.core as mx
j1 = mx.ones(100_000)
j2 = mx.ones(100_000)
j3 = 2 * mx.ones(100_000)
m1 = mx.zeros(100_000)
m2 = mx.zeros(100_000)
m3 = mx.zeros(100_000)
result = wigner_3j(j1, j2, j3, m1, m2, m3)  # all on GPU

# 6j and 9j symbols
s6 = wigner_6j(1, 1, 1, 1, 1, 1)
s9 = wigner_9j(1, 1, 0, 1, 1, 0, 0, 0, 0)

API

Function	Signature	Description
wigner_3j	(j1, j2, j3, m1, m2, m3)	Wigner 3j symbol
wigner_6j	(j1, j2, j3, j4, j5, j6)	Wigner 6j symbol (Racah W)
wigner_9j	(j1, ..., j9)	Wigner 9j symbol
clebsch_gordan	(j1, m1, j2, m2, J, M)	Clebsch-Gordan coefficient

All functions accept scalars or MLX arrays. Array inputs are batch-computed on GPU.

Benchmark

python benchmark.py

Typical results on M3 Max (vs sympy on CPU):

Batch size	mlx-wigner (GPU)	sympy (CPU)	Speedup
1	~0.5 ms	~2 ms	4x
1,000	~1 ms	~2,000 ms	2,000x
100,000	~10 ms	—	—

The GPU advantage grows with batch size due to massive parallelism.

Algorithm

• Wigner 3j: Racah formula with log-gamma (Lanczos approximation) for numerical stability
• Wigner 6j: Direct Racah sum formula (not built from 3j products)
• Wigner 9j: Sum over 6j triple products with internal summation variable
• Selection rules: Vectorized boolean masks applied before computation

Half-integer angular momenta are fully supported.

Testing

pip install -e ".[dev]"
pytest tests/ -v

Tests validate against exact values from sympy.physics.wigner.

License

MIT License. Copyright (c) 2026 Sheng-Kai Huang.