mlx-hyp2f1 0.1.0

GPU-accelerated hypergeometric functions (2F1, 1F1, 0F1) for Apple MLX

mlx-hyp2f1

GPU-accelerated hypergeometric functions for Apple MLX.

JAX has hyp2f1 but PyTorch, CuPy, and MLX do not. This package fills the gap for MLX, providing vectorized implementations that run on Apple Silicon GPU.

Features

• hyp2f1(a, b, c, z) -- Gauss hypergeometric function 2F1
• hyp1f1(a, b, z) -- Confluent (Kummer) hypergeometric function 1F1
• hyp0f1(b, z) -- Bessel-related hypergeometric function 0F1
• Fully vectorized: compute over arrays of parameters simultaneously on GPU
• Numerically stable via log-gamma Pochhammer symbols
• Analytic continuation for |z| >= 1 using linear transformation formulas
• Pure MLX -- no SciPy runtime dependency

Installation

pip install mlx-hyp2f1

Or from source:

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

Usage

import mlx.core as mx
from mlx_hyp2f1 import hyp2f1, hyp1f1, hyp0f1

# Scalar
result = hyp2f1(0.5, 1.0, 1.5, 0.3)

# Vectorized over z
z = mx.linspace(0.0, 0.95, 100)
result = hyp2f1(0.5, 1.0, 1.5, z)

# Vectorized over all parameters
a = mx.array([0.5, 1.0, 2.0])
b = mx.array([1.0, 1.5, 0.5])
c = mx.array([1.5, 2.0, 3.0])
z = mx.array([0.3, 0.5, 0.8])
result = hyp2f1(a, b, c, z)

# Confluent hypergeometric (Kummer)
result = hyp1f1(1.0, 2.0, mx.linspace(-5.0, 5.0, 100))

# Bessel-related
result = hyp0f1(1.0, mx.linspace(-10.0, 10.0, 100))

Cosmological growth factor D(a)

import mlx.core as mx
from mlx_hyp2f1 import hyp2f1

def growth_factor(a, omega_m=0.3):
    """Linear growth factor D(a) for flat LCDM."""
    omega_l = 1.0 - omega_m
    x = -omega_l / omega_m * a**3
    return a * hyp2f1(1.0/3.0, 1.0, 11.0/6.0, x)

Benchmark

python -m mlx_hyp2f1.benchmark

Testing

pip install -e ".[dev]"
pytest

Applications

• Cosmological growth factor D(a)
• Beta distribution CDF / PDF
• Angular momentum coupling coefficients
• Legendre / Jacobi / Gegenbauer polynomials
• Scattering amplitudes in quantum mechanics

License

MIT -- Sheng-Kai Huang