mlx-airy 0.1.1

GPU-accelerated Airy functions for Apple Silicon using MLX

mlx-airy

GPU-accelerated Airy functions for Apple Silicon using MLX.

Neither MLX nor JAX currently provide Airy functions (see JAX issue #25244). This package fills that gap for the MLX ecosystem.

What are Airy functions?

Airy functions Ai(x) and Bi(x) are the two linearly independent solutions to the Airy differential equation:

y'' - x y = 0

They appear throughout physics and engineering:

• Quantum mechanics -- WKB connection formulas, tunneling through linear potentials
• Optics -- diffraction near caustics, rainbow scattering
• Fluid dynamics -- boundary layers, Stokes phenomena
• Astrophysics -- gravitational lensing near fold caustics

Installation

pip install -e ".[test,benchmark]"

Requires Python >= 3.10 and Apple Silicon (MLX).

Usage

import mlx.core as mx
from mlx_airy import airy

x = mx.linspace(-15.0, 15.0, 10000)
ai, aip, bi, bip = airy(x)  # Ai(x), Ai'(x), Bi(x), Bi'(x)

All four functions are returned in a single call, matching the scipy.special.airy convention.

Implementation

Region	Method	Reference
|x| < 5	Taylor series (50 terms, Horner)	DLMF 9.4
x >= 5	Asymptotic expansion	DLMF 9.7.5--9.7.8
x <= -5	Oscillatory asymptotic	DLMF 9.7.9--9.7.12

Computation is float32 throughout, running entirely on the Apple GPU.

Tests

pytest tests/ -v

Benchmark

python -m mlx_airy.benchmark

Compares accuracy and throughput against scipy.special.airy.

Measured on Apple M1 Max (MLX 0.31.1, SciPy 1.16.2, float32), x in [-15, 15], all four outputs (Ai, Ai', Bi, Bi'):

N	scipy (ms)	mlx (ms)	Speedup
1,000	3.30	32.07	0.1x
10,000	15.76	31.65	0.5x
100,000	147.69	35.55	4.2x
1,000,000	1393.08	209.09	6.7x

mlx-airy wins above ~30-50k points, where the Metal dispatch cost is amortised; below that, scipy's vectorised C path is faster. Worst-case float32 relative error is ~3e-4. Full numbers in benchmark_results.md.

Author

Sheng-Kai Huang ([e-mail])

License

MIT