mlx-airy 0.1.0

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.

Author

Sheng-Kai Huang ([e-mail])

License

MIT