# mlx-gamma

GPU-accelerated gamma and lgamma functions for [Apple MLX](https://github.com/ml-explore/mlx).

MLX is the only major ML framework missing `gamma`/`lgamma` support
([mlx#2030](https://github.com/apple/mlx/issues/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

```bash
pip install mlx-gamma
```

Or from source:

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

## Quick start

```python
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

```bash
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

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

## License

MIT -- see [LICENSE](LICENSE).

## Author

Sheng-Kai Huang (akai@fawstudio.com)

## References

- MLX issue: [apple/mlx#2030](https://github.com/apple/mlx/issues/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.
