Spherical Bessel functions on MLX: 570x speedup over scipy

[akaiHuang]

I implemented GPU-accelerated spherical Bessel functions using piecewise Chebyshev interpolation on MLX.

Performance (M1 Max)

Size (ells × points)	scipy (CPU)	MLX (GPU eval)	Speedup
100 × 1000	0.31s	0.003s	96x
500 × 5000	8.95s	0.016s	557x
525 × 10000	19.6s	0.034s	579x

Table build is a one-time cost (~5s for 500 ells). Subsequent evaluations use cached GPU coefficients.

Accuracy

Metric	Value
Max absolute error	6.3 × 10⁻⁷
Median relative error	0.005%
P99 relative error	0.24%

Method

Piecewise Chebyshev interpolation: precompute j_l(x) at Chebyshev nodes on CPU, store coefficients on GPU, evaluate via matrix multiply. Supports both j_l(x) and j_l'(x).

Use case

Spherical Bessel functions are needed in CMB cosmology (Boltzmann solvers), acoustics, and electromagnetic simulations. No GPU framework currently provides them natively.

Code

• GitHub: https://github.com/akaiHuang/mlx-bessel
• License: MIT

from mlx_bessel import BesselTable
import mlx.core as mx

table = BesselTable(ell_max=500, x_max=5000)
x = mx.linspace(0.1, 5000, 10000)
jl = table.eval_jl(x)  # shape: (501, 10000), on GPU

Would there be interest in adding spherical Bessel functions to mlx or mlx-examples?