mlx-gamma 0.1.0

GPU-accelerated gamma and lgamma functions for Apple MLX

mlx-gamma

GPU-accelerated gamma and lgamma functions for Apple MLX.

MLX is the only major ML framework missing gamma/lgamma support (mlx#2030). PyTorch, JAX, and CuPy all provide these. mlx-gamma fills that gap with pure-MLX vectorized implementations that run entirely on the Apple GPU.

Installation

pip install mlx-gamma

Or from source:

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

Quick start

import mlx.core as mx
from mlx_gamma import lgamma, gamma, digamma, beta

x = mx.array([1.0, 2.0, 3.0, 5.0, 10.0, 100.0])

lgamma(x)   # log-gamma
gamma(x)    # gamma function (with sign handling)
digamma(x)  # psi function, d/dx ln Gamma(x)

beta(mx.array([2.0, 3.0]), mx.array([3.0, 4.0]))  # Beta function

Functions

Function	Description	Domain
lgamma(x)	Log of absolute value of Gamma(x)	x != 0, -1, -2, ...
gamma(x)	Gamma function via exp(lgamma(x)) with sign	x != 0, -1, -2, ...
digamma(x)	Psi function (logarithmic derivative of Gamma)	x != 0, -1, -2, ...
beta(a, b)	Beta function B(a,b) = Gamma(a)Gamma(b)/Gamma(a+b)	a, b > 0

Algorithm details

• lgamma: Lanczos approximation (g=7, 9 coefficients) for x >= 0.5; reflection formula for x < 0.5. Accurate to ~13 decimal digits in float64, ~6 in float32.
• gamma: exp(lgamma(x)) with sign correction for negative arguments.
• digamma: Asymptotic expansion for x >= 7; recurrence relation to shift small arguments upward; reflection formula for x < 0.
• beta: Computed via exp(lgamma(a) + lgamma(b) - lgamma(a+b)) for numerical stability.

Benchmarks

python benchmark.py

Compares against scipy.special.gammaln on CPU. Typical results on M1 Max for 1M elements:

Function	scipy (CPU)	mlx-gamma (GPU)	Speedup
lgamma	~12 ms	~0.8 ms	~15x

Running tests

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

License

MIT -- see LICENSE.

Author

Sheng-Kai Huang ([e-mail])

References

• MLX issue: apple/mlx#2030
• Lanczos, C. (1964). "A Precision Approximation of the Gamma Function." SIAM J. Numer. Anal. 1: 86--96.
• Pugh, G. R. (2004). "An Analysis of the Lanczos Gamma Approximation." PhD thesis, University of British Columbia.