From d8c067ea560c669521c4b7ffce0fd5457b087777 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 11 Aug 2026 21:49:28 -0500 Subject: [PATCH 1/4] Add MLX dispatch for gammainc and gammaincc --- pytensor/link/mlx/dispatch/scalar/__init__.py | 10 +- pytensor/link/mlx/dispatch/scalar/gammainc.py | 371 ++++++++++++++++++ tests/link/mlx/scalar/test_gammainc.py | 217 ++++++++++ 3 files changed, 597 insertions(+), 1 deletion(-) create mode 100644 pytensor/link/mlx/dispatch/scalar/gammainc.py create mode 100644 tests/link/mlx/scalar/test_gammainc.py diff --git a/pytensor/link/mlx/dispatch/scalar/__init__.py b/pytensor/link/mlx/dispatch/scalar/__init__.py index 9f29f83c48..12b9dc7cec 100644 --- a/pytensor/link/mlx/dispatch/scalar/__init__.py +++ b/pytensor/link/mlx/dispatch/scalar/__init__.py @@ -1 +1,9 @@ -from pytensor.link.mlx.dispatch.scalar import basic, bessel, erf, erfcinv, gamma, math +from pytensor.link.mlx.dispatch.scalar import ( + basic, + bessel, + erf, + erfcinv, + gamma, + gammainc, + math, +) diff --git a/pytensor/link/mlx/dispatch/scalar/gammainc.py b/pytensor/link/mlx/dispatch/scalar/gammainc.py new file mode 100644 index 0000000000..13db5e985a --- /dev/null +++ b/pytensor/link/mlx/dispatch/scalar/gammainc.py @@ -0,0 +1,371 @@ +import mlx.core as mx +import numpy as np + +from pytensor.link.mlx.dispatch.basic import mlx_funcify +from pytensor.link.mlx.dispatch.scalar.helpers import _exp, _working_precision +from pytensor.link.mlx.dispatch.scalar.math import _lanczos_log_gamma +from pytensor.scalar.math import Erfc, GammaInc, GammaIncC + + +# The regularized incomplete gamma functions have no single expansion that holds over +# the whole domain, so three cover it between them: +# +# the ascending series for P below x = a + 1, +# the Lentz continued fraction for Q above it, +# and Temme's uniform asymptotic expansion (DLMF 8.12) through the transition at +# large a, where the other two need a number of terms growing with sqrt(a). +# +# The trip counts below are fixed because a is a runtime tensor -- a shape parameter, +# usually a random variable -- so nothing about it is known when the graph is built. +# Temme is what makes that affordable: its cost does not depend on a at all, and its +# accuracy improves as a grows, which is exactly where the other two fall apart. +_SERIES_TERMS = 50 +_CF_TERMS = 50 +# The window is set by float32 rather than float64: the series and the continued +# fraction both reach their answer through a log prefactor that subtracts two numbers of +# order a log a, so single precision loses digits in proportion to a, while Temme's form +# has no such subtraction. Widening past this trades float64 accuracy for no further +# float32 gain +_TEMME_MIN_A = 15.0 +_TEMME_MIN_RATIO = 0.4 +_TEMME_MAX_RATIO = 1.8 + +# 0.25**25 sits below float64 epsilon, so the series is exact to working precision +# everywhere it is selected +_ETA_SERIES_TERMS = 25 +_ETA_SERIES_MAX_SIGMA = 0.25 +# The continued fraction's zero-denominator floor. It has to be a normal at float32 as +# well as float64, or it rounds to zero there and takes both of Lentz's guards with it +_TINY = 1e-30 + +# d[k][n] from DLMF 8.12.12, generated by the script scipy ships at +# scipy/special/_precompute/gammainc_asy.py. Ten k by fourteen n holds 1e-13 from +# a = 15 upward across the window this uses it in; scipy's own table is larger because +# it applies the expansion further from x = a than we need to. +_TEMME_D = ( + ( + -0.3333333333333333, + 0.08333333333333333, + -0.014814814814814815, + 0.0011574074074074073, + 0.0003527336860670194, + -0.0001787551440329218, + 3.919263178522438e-05, + -2.185448510679992e-06, + -1.85406221071516e-06, + 8.296711340953087e-07, + -1.7665952736826078e-07, + 6.707853543401498e-09, + 1.0261809784240309e-08, + -4.382036018453353e-09, + ), + ( + -0.001851851851851852, + -0.003472222222222222, + 0.0026455026455026454, + -0.0009902263374485596, + 0.00020576131687242798, + -4.018775720164609e-07, + -1.8098550334489977e-05, + 7.64916091608111e-06, + -1.6120900894563446e-06, + 4.647127802807434e-09, + 1.378633446915721e-07, + -5.752545603517705e-08, + 1.1951628599778148e-08, + -1.7543241719747647e-11, + ), + ( + 0.004133597883597883, + -0.0026813271604938273, + 0.0007716049382716049, + 2.0093878600823047e-06, + -0.0001073665322636516, + 5.2923448829120125e-05, + -1.2760635188618728e-05, + 3.423578734096138e-08, + 1.3721957309062934e-06, + -6.298992138380055e-07, + 1.4280614206064242e-07, + -2.0477098421990866e-10, + -1.409252991086752e-08, + 6.228974084922022e-09, + ), + ( + 0.0006494341563786008, + 0.00022947209362139917, + -0.0004691894943952557, + 0.00026772063206283885, + -7.561801671883977e-05, + -2.396505113867297e-07, + 1.1082654115347302e-05, + -5.6749528269915965e-06, + 1.4230900732435883e-06, + -2.7861080291528143e-11, + -1.6958404091930278e-07, + 8.099464905388083e-08, + -1.9111168485973655e-08, + 2.3928620439808118e-12, + ), + ( + -0.0008618882909167117, + 0.0007840392217200666, + -0.0002990724803031902, + -1.4638452578843418e-06, + 6.641498215465122e-05, + -3.968365047179435e-05, + 1.1375726970678419e-05, + 2.507497226237533e-10, + -1.6954149536558305e-06, + 8.907507532205309e-07, + -2.292934834000805e-07, + 2.956794137544049e-11, + 2.8865829742708783e-08, + -1.4189739437803219e-08, + ), + ( + -0.00033679855336635813, + -6.972813758365857e-05, + 0.0002772753244959392, + -0.00019932570516188847, + 6.797780477937208e-05, + 1.419062920643967e-07, + -1.3594048189768693e-05, + 8.018470256334202e-06, + -2.291481176508095e-06, + -3.252473551298454e-10, + 3.4652846491085265e-07, + -1.8447187191171344e-07, + 4.8240967037894184e-08, + -1.7989466721743514e-14, + ), + ( + 0.0005313079364639922, + -0.0005921664373536939, + 0.0002708782096718045, + 7.902353232660328e-07, + -8.153969367561969e-05, + 5.61168275310625e-05, + -1.8329116582843375e-05, + -3.0796134506033047e-09, + 3.465155368803609e-06, + -2.0291327396058603e-06, + 5.788792863149004e-07, + 2.338630673826657e-13, + -8.828600746330484e-08, + 4.7435958880408125e-08, + ), + ( + 0.00034436760689237765, + 5.171790908260592e-05, + -0.00033493161081142234, + 0.0002812695154763237, + -0.00010976582244684731, + -1.2741009095484485e-07, + 2.7744451511563645e-05, + -1.8263488805711332e-05, + 5.7876949497350525e-06, + 4.93875893393627e-10, + -1.0595367014026043e-06, + 6.166714376110408e-07, + -1.7562973359060463e-07, + -1.2974473287015437e-12, + ), + ( + -0.0006526239185953094, + 0.0008394987206720873, + -0.000438297098541721, + -6.969091458420552e-07, + 0.00016644846642067547, + -0.00012783517679769218, + 4.629953263691304e-05, + 4.557909867922708e-09, + -1.0595271125805195e-05, + 6.783342904865167e-06, + -2.1075476666258803e-06, + -1.7213731432817144e-11, + 3.773587741611098e-07, + -2.1867506700122867e-07, + ), + ( + -0.0005967612901927463, + -7.204895416020011e-05, + 0.0006782308837667328, + -0.0006401475260262758, + 0.00027750107634328704, + 1.819700838046515e-07, + -8.479507117068503e-05, + 6.105192082501531e-05, + -2.1073920183404862e-05, + -8.858589014125599e-10, + 4.5284535953805374e-06, + -2.8427815022504407e-06, + 8.708234177864641e-07, + 3.688610187171431e-12, + ), +) + + +def _temme_eta(a, x, const): + r"""Temme's variable, with :math:`\eta^2 / 2 = \lambda - 1 - \log \lambda`. + + Even in the ``log1p`` form the two terms cancel to the size of their own difference + as :math:`x \to a`, which is the transition Temme exists to cover. The Maclaurin + series of that difference has no subtraction in it at all, so it carries the + small-:math:`\sigma` end. The loss is amplified downstream: :math:`\eta` enters an + exponent scaled by :math:`a`, so a relative error in it comes back multiplied by + :math:`a \eta^2 / 2`. + """ + sigma = (x - a) / a + + series = const(0.0) + for k in reversed(range(_ETA_SERIES_TERMS)): + series = series * (-sigma) + const(2.0 / (k + 2.0)) + value = mx.where( + mx.abs(sigma) <= const(_ETA_SERIES_MAX_SIGMA), + sigma * sigma * series, + const(-2.0) * (mx.log1p(sigma) - sigma), + ) + return mx.sign(sigma) * mx.sqrt(mx.maximum(value, const(0.0))) + + +def _temme(a, x, sign, const, erfc): + r"""One tail by the uniform asymptotic expansion, for large :math:`a` near the + transition. + + Parameters + ---------- + sign : float + ``+1.0`` for :math:`Q(a, x)`, ``-1.0`` for :math:`P(a, x)`. Both tails come out + of the same expansion, by + + .. math:: + \operatorname{erfc}(z) + \operatorname{erfc}(-z) = 2 + + so each is built from terms of its own magnitude rather than as a subtraction + from one. + """ + eta = _temme_eta(a, x, const) + + total = const(0.0) + a_power = const(1.0) + for row in _TEMME_D: + coefficient = const(row[-1]) + for value in reversed(row[:-1]): + coefficient = coefficient * eta + const(value) + total = total + coefficient * a_power + a_power = a_power / a + + correction = ( + _exp(const(-0.5) * a * eta * eta, const) + * mx.rsqrt(const(2.0 * np.pi) * a) + * total + ) + tail = const(0.5) * erfc(const(sign) * eta * mx.sqrt(a * const(0.5))) + return tail + const(sign) * correction + + +def _gammainc_series(a, x, const): + """``P(a, x)`` by the ascending series, for ``x`` at or below ``a + 1``.""" + term = const(1.0) + total = const(1.0) + for n in range(1, _SERIES_TERMS): + term = term * x / (a + const(float(n))) + total = total + term + + # the whole x**a exp(-x) / Gamma(a + 1) prefactor in log space, so that a large a + # does not overflow on its way to a small answer. x = 0 is substituted rather than + # clamped: clamping puts a * log(tiny) in the exponent, which for a small a lands on + # a finite nonzero answer instead of the exact zero P(a, 0) has + positive = x > const(0.0) + log_prefactor = ( + -x + + a * mx.log(mx.where(positive, x, const(1.0))) + - _lanczos_log_gamma(a, const) + ) + # x - x rather than a zero constant, so a nan argument still propagates + return mx.where(positive, total * _exp(log_prefactor, const), x - x) + + +def _gammaincc_cf(a, x, const): + """``Q(a, x)`` by the Lentz continued fraction, for ``x`` above ``a + 1``.""" + tiny = const(_TINY) + b = x + const(1.0) - a + c = const(1.0) / tiny + d = const(1.0) / mx.where(b == const(0.0), tiny, b) + h = d + + for i in range(1, _CF_TERMS): + an = -const(float(i)) * (const(float(i)) - a) + b = b + const(2.0) + d = an * d + b + d = mx.where(mx.abs(d) < tiny, tiny, d) + c = b + an / c + c = mx.where(mx.abs(c) < tiny, tiny, c) + d = const(1.0) / d + h = h * d * c + + # the x**a exp(-x) / Gamma(a) prefactor; _lanczos_log_gamma is shifted by one, so + # Gamma(a) is its value at a - 1 + log_prefactor = -x + a * mx.log(x) - _lanczos_log_gamma(a - const(1.0), const) + return h * _exp(log_prefactor, const) + + +def _incomplete_gamma(a, x, const, erfc, *, lower): + r"""One regularized incomplete gamma tail over the whole domain. + + Three expansions cover the plane, chosen per element. Whichever tail is asked for is + always the one the selected expansion produces at full precision: Temme supplies + either directly, the series owns the region where :math:`Q` is order one, and the + continued fraction owns the region where :math:`P` is, so the complement is never a + subtraction of two nearby numbers. + + Parameters + ---------- + lower : bool + Return :math:`P(a, x)` when true, :math:`Q(a, x)` when false. + """ + use_temme = ( + (a >= const(_TEMME_MIN_A)) + & (x >= const(_TEMME_MIN_RATIO) * a) + & (x <= const(_TEMME_MAX_RATIO) * a) + ) + use_series = ~use_temme & (x <= a + const(1.0)) + + # Every branch runs for every element, so each is handed an argument it can survive + # rather than the one belonging to whichever branch will actually be selected + temme = _temme(a, mx.where(use_temme, x, a), -1.0 if lower else 1.0, const, erfc) + series = _gammainc_series(a, mx.where(use_series, x, a * const(0.25)), const) + fraction = _gammaincc_cf( + a, mx.where(use_temme | use_series, a + const(2.0), x), const + ) + if lower: + fraction = const(1.0) - fraction + else: + series = const(1.0) - series + + return mx.where(use_temme, temme, mx.where(use_series, series, fraction)) + + +@mlx_funcify.register(GammaIncC) +def mlx_funcify_GammaIncC(op, **kwargs): + erfc = mlx_funcify(Erfc()) + + def gammaincc(a, x): + z, const, out_dtype = _working_precision(x) + shape = mx.array(a).astype(z.dtype) + return _incomplete_gamma(shape, z, const, erfc, lower=False).astype(out_dtype) + + return gammaincc + + +@mlx_funcify.register(GammaInc) +def mlx_funcify_GammaInc(op, **kwargs): + erfc = mlx_funcify(Erfc()) + + def gammainc(a, x): + z, const, out_dtype = _working_precision(x) + shape = mx.array(a).astype(z.dtype) + return _incomplete_gamma(shape, z, const, erfc, lower=True).astype(out_dtype) + + return gammainc diff --git a/tests/link/mlx/scalar/test_gammainc.py b/tests/link/mlx/scalar/test_gammainc.py new file mode 100644 index 0000000000..2b061cf105 --- /dev/null +++ b/tests/link/mlx/scalar/test_gammainc.py @@ -0,0 +1,217 @@ +from functools import partial + +import numpy as np +import pytest +import scipy.special + +import pytensor.tensor as pt +from pytensor.scalar.math import GammaInc, GammaIncC +from pytensor.tensor.type import vector +from tests.link.mlx.test_basic import compare_mlx_and_py + + +mlx = pytest.importorskip("mlx.core") +from pytensor.link.mlx.dispatch import mlx_funcify +from pytensor.link.mlx.dispatch.scalar.gammainc import ( + _TEMME_MAX_RATIO, + _TEMME_MIN_A, + _TEMME_MIN_RATIO, + _TINY, +) + + +def _sample(rng, shape_range, ratio_range, n): + """Draw log-spaced ``(a, x)`` with ``x`` set as a multiple of ``a``. + + The expansion is chosen on ``x / a``, so a region is only reached by sampling that + ratio rather than ``x`` on its own. + """ + a = np.exp(rng.uniform(*np.log(shape_range), n)) + return a, a * np.exp(rng.uniform(*np.log(ratio_range), n)) + + +# Each region selects a different expansion: Temme needs a >= 15 with x within [0.4a, +# 1.8a], the ascending series takes what is left below x = a + 1, the continued fraction +# what is left above it. float32 loosens where the series and the fraction reach their +# answer through a log prefactor that subtracts two numbers of order a log a. +@pytest.mark.parametrize("dtype", ["float32", "float64"], ids=str) +@pytest.mark.parametrize( + "shape_range, ratio_range, tolerances", + [ + ((1e-2, 5.0), (1e-2, 0.9), {"float32": 1e-4, "float64": 1e-12}), + ((20.0, 5e4), (0.4, 1.8), {"float32": 3e-4, "float64": 1e-10}), + ((1e-2, 10.0), (5.0, 50.0), {"float32": 1e-3, "float64": 1e-10}), + ((50.0, 5e3), (0.05, 0.35), {"float32": 5e-4, "float64": 1e-11}), + ], + ids=["series", "temme", "continued-fraction", "series-far-left"], +) +@pytest.mark.parametrize( + "op", [pt.gammainc, pt.gammaincc], ids=["gammainc", "gammaincc"] +) +def test_incomplete_gamma(op, shape_range, ratio_range, tolerances, dtype): + a = vector("a", dtype=dtype) + x = vector("x", dtype=dtype) + a_test_value, x_test_value = _sample( + np.random.default_rng(11), shape_range, ratio_range, 201 + ) + + compare_mlx_and_py( + [a, x], + [op(a, x)], + [a_test_value.astype(dtype), x_test_value.astype(dtype)], + # Far enough into the tail the answer is a subnormal, where one ULP is percents + # and no relative claim means anything; atol retires those elements at the + # dtype's smallest normal rather than dropping them from the sample + assert_fn=partial( + np.testing.assert_allclose, + rtol=tolerances[dtype], + atol=np.finfo(dtype).tiny, + ), + ) + + +@pytest.mark.parametrize( + "shape_range, ratio_range", + [ + ((1e-2, 5.0), (1e-2, 0.9)), + ((20.0, 5e4), (0.4, 1.8)), + ((1e-2, 10.0), (5.0, 50.0)), + ((50.0, 5e3), (0.05, 0.35)), + ], + ids=["series", "temme", "continued-fraction", "series-far-left"], +) +def test_incomplete_gamma_float64_precision(shape_range, ratio_range): + a_test_value, x_test_value = _sample( + np.random.default_rng(23), shape_range, ratio_range, 2001 + ) + + with mlx.stream(mlx.cpu): + a_mlx = mlx.array(a_test_value, dtype=mlx.float64) + x_mlx = mlx.array(x_test_value, dtype=mlx.float64) + lower = np.asarray(mlx_funcify(GammaInc())(a_mlx, x_mlx)) + upper = np.asarray(mlx_funcify(GammaIncC())(a_mlx, x_mlx)) + + # Past 1e-290 a tail has run out of exponent rather than of accuracy, so those + # elements are dropped instead of carrying a relative claim about a denormal + truth_lower = scipy.special.gammainc(a_test_value, x_test_value) + truth_upper = scipy.special.gammaincc(a_test_value, x_test_value) + for got, truth in ((lower, truth_lower), (upper, truth_upper)): + keep = truth > 1e-290 + np.testing.assert_allclose(got[keep], truth[keep], rtol=1e-10) + + +def test_incomplete_gamma_small_tail(): + # The regression this exists for: taking one tail as one minus the other cancels + # everything once that tail runs below the other's rounding. At a = 200, x = 0.5a + # the true P is 9.3e-19, so a complement would report zero with 100% error. Both + # tails come out of the same expansion directly instead. + a_test_value = np.array([200.0, 200.0, 2000.0, 2000.0, 5000.0]) + x_test_value = a_test_value * np.array([0.5, 0.8, 0.8, 0.9, 0.85]) + + with mlx.stream(mlx.cpu): + lower = np.asarray( + mlx_funcify(GammaInc())( + mlx.array(a_test_value, dtype=mlx.float64), + mlx.array(x_test_value, dtype=mlx.float64), + ) + ) + + truth = scipy.special.gammainc(a_test_value, x_test_value) + assert truth.max() < 1e-2, "these points have to sit in the tail to test anything" + np.testing.assert_allclose(lower, truth, rtol=1e-10) + + +def test_incomplete_gamma_branch_boundaries(): + # The three expansions meet along a = 15, x = 0.4a, x = 1.8a and x = a + 1, and the + # answer has to agree from both sides of every seam. A wrong coefficient or trip count + # in one expansion shows up here as a step across the seam it owns, which sampling + # inside a region cannot see. The thresholds are imported rather than written out, so + # moving a boundary moves the test with it. + offsets = np.array([-1e-6, -1e-9, 0.0, 1e-9, 1e-6]) + shapes, values = [], [] + for shape in (5.0, _TEMME_MIN_A - 1e-6, _TEMME_MIN_A, _TEMME_MIN_A + 1e-6, 3000.0): + for ratio in (_TEMME_MIN_RATIO, _TEMME_MAX_RATIO): + shapes.append(np.full(offsets.size, shape)) + values.append(shape * (ratio + offsets)) + shapes.append(np.full(offsets.size, shape)) + values.append(shape + 1.0 + offsets * shape) + a_test_value = np.concatenate(shapes) + x_test_value = np.concatenate(values) + + with mlx.stream(mlx.cpu): + a_mlx = mlx.array(a_test_value, dtype=mlx.float64) + x_mlx = mlx.array(x_test_value, dtype=mlx.float64) + lower = np.asarray(mlx_funcify(GammaInc())(a_mlx, x_mlx)) + upper = np.asarray(mlx_funcify(GammaIncC())(a_mlx, x_mlx)) + + np.testing.assert_allclose( + lower, scipy.special.gammainc(a_test_value, x_test_value), rtol=1e-10 + ) + np.testing.assert_allclose( + upper, scipy.special.gammaincc(a_test_value, x_test_value), rtol=1e-10 + ) + + +def test_continued_fraction_floor_survives_float32(): + # The continued fraction seeds one Lentz variable at 1 / _TINY and guards both against + # it. A floor below float32's smallest normal -- 1e-300 was -- rounds to zero there, + # which seeds inf and turns each guard into |d| < 0, so neither can ever fire. + assert np.float32(_TINY) > 0.0 + assert np.isfinite(np.float32(1.0) / np.float32(_TINY)) + + +def test_incomplete_gamma_complement(): + # P + Q == 1 has to hold across every expansion and both sides of each boundary, + # including where one tail is hundreds of orders below the other + a_test_value, x_test_value = _sample( + np.random.default_rng(37), (1e-2, 1e5), (1e-3, 50.0), 4001 + ) + + with mlx.stream(mlx.cpu): + a_mlx = mlx.array(a_test_value, dtype=mlx.float64) + x_mlx = mlx.array(x_test_value, dtype=mlx.float64) + total = np.asarray(mlx_funcify(GammaInc())(a_mlx, x_mlx)) + np.asarray( + mlx_funcify(GammaIncC())(a_mlx, x_mlx) + ) + + np.testing.assert_allclose(total, 1.0, rtol=1e-12) + + +@pytest.mark.parametrize( + "op", [pt.gammainc, pt.gammaincc], ids=["gammainc", "gammaincc"] +) +def test_incomplete_gamma_edge_cases(op): + # x = 0 is exactly P = 0, which a log prefactor clamped away from zero rather than + # substituted gets wrong by a whole answer: a * log(tiny) in the exponent lands on + # 1e-150 for a = 0.5. The rest pin the saturated ends and nan propagation. + a = vector("a", dtype="float64") + x = vector("x", dtype="float64") + a_test_value = np.array([1.0, 0.5, 1e-3, 200.0, 1.0, 200.0, 1.0, 1.0]) + x_test_value = np.array([0.0, 0.0, 0.0, 0.0, 1e30, 1e30, 1e-8, np.nan]) + + compare_mlx_and_py([a, x], [op(a, x)], [a_test_value, x_test_value]) + + +@pytest.mark.parametrize( + "dtype, rtol", [("float32", 1e-4), ("float64", 1e-6)], ids=["float32", "float64"] +) +def test_incomplete_gamma_grad(dtype, rtol): + # Only the gradient in the value has a dispatch: the one in the shape parameter is a + # ScalarLoop, which the MLX backend does not convert yet. + # + # The gradient graph is pytensor's own exp(...) rather than anything in this module, + # so it inherits mx.exp, a float32 kernel at every dtype that flushes to zero below + # exp(-90). That sets the float64 tolerance, and it is why the sampled range keeps + # the density above that floor rather than running out into the tail. + a = vector("a", dtype=dtype) + x = vector("x", dtype=dtype) + a_test_value, x_test_value = _sample( + np.random.default_rng(41), (0.5, 20.0), (0.2, 3.0), 101 + ) + + compare_mlx_and_py( + [a, x], + [pt.grad(pt.gammainc(a, x).sum(), x), pt.grad(pt.gammaincc(a, x).sum(), x)], + [a_test_value.astype(dtype), x_test_value.astype(dtype)], + assert_fn=partial(np.testing.assert_allclose, rtol=rtol), + ) From b5422a4c2a92bb00089a69d4335c780e35c36499 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 12 Aug 2026 20:10:32 -0500 Subject: [PATCH 2/4] Saturate the MLX incomplete gamma at an infinite argument --- pytensor/link/mlx/dispatch/scalar/gammainc.py | 9 +++++++-- tests/link/mlx/scalar/test_gammainc.py | 10 +++++++--- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/pytensor/link/mlx/dispatch/scalar/gammainc.py b/pytensor/link/mlx/dispatch/scalar/gammainc.py index 13db5e985a..7c40f92cf4 100644 --- a/pytensor/link/mlx/dispatch/scalar/gammainc.py +++ b/pytensor/link/mlx/dispatch/scalar/gammainc.py @@ -331,20 +331,25 @@ def _incomplete_gamma(a, x, const, erfc, *, lower): & (x <= const(_TEMME_MAX_RATIO) * a) ) use_series = ~use_temme & (x <= a + const(1.0)) + # x = inf saturates instead of reaching the continued fraction, where b = x + 1 - a is + # itself inf and the Lentz step forms (1 / inf) * inf. Tested rather than compared + # against a constant, so no infinity is ever materialized for mx.compile to inline + saturated = mx.isinf(x) & (x > const(0.0)) # Every branch runs for every element, so each is handed an argument it can survive # rather than the one belonging to whichever branch will actually be selected temme = _temme(a, mx.where(use_temme, x, a), -1.0 if lower else 1.0, const, erfc) series = _gammainc_series(a, mx.where(use_series, x, a * const(0.25)), const) fraction = _gammaincc_cf( - a, mx.where(use_temme | use_series, a + const(2.0), x), const + a, mx.where(use_temme | use_series | saturated, a + const(2.0), x), const ) if lower: fraction = const(1.0) - fraction else: series = const(1.0) - series - return mx.where(use_temme, temme, mx.where(use_series, series, fraction)) + selected = mx.where(use_temme, temme, mx.where(use_series, series, fraction)) + return mx.where(saturated, const(1.0 if lower else 0.0), selected) @mlx_funcify.register(GammaIncC) diff --git a/tests/link/mlx/scalar/test_gammainc.py b/tests/link/mlx/scalar/test_gammainc.py index 2b061cf105..e1a71faed3 100644 --- a/tests/link/mlx/scalar/test_gammainc.py +++ b/tests/link/mlx/scalar/test_gammainc.py @@ -183,11 +183,15 @@ def test_incomplete_gamma_complement(): def test_incomplete_gamma_edge_cases(op): # x = 0 is exactly P = 0, which a log prefactor clamped away from zero rather than # substituted gets wrong by a whole answer: a * log(tiny) in the exponent lands on - # 1e-150 for a = 0.5. The rest pin the saturated ends and nan propagation. + # 1e-150 for a = 0.5. x = inf saturates, where the continued fraction it would + # otherwise reach forms (1 / inf) * inf and returns nan. The rest pin the large-but- + # finite end and nan propagation. a = vector("a", dtype="float64") x = vector("x", dtype="float64") - a_test_value = np.array([1.0, 0.5, 1e-3, 200.0, 1.0, 200.0, 1.0, 1.0]) - x_test_value = np.array([0.0, 0.0, 0.0, 0.0, 1e30, 1e30, 1e-8, np.nan]) + a_test_value = np.array([1.0, 0.5, 1e-3, 200.0, 1.0, 200.0, 1.0, 0.5, 200.0, 1.0]) + x_test_value = np.array( + [0.0, 0.0, 0.0, 0.0, 1e30, 1e30, np.inf, np.inf, np.inf, np.nan] + ) compare_mlx_and_py([a, x], [op(a, x)], [a_test_value, x_test_value]) From 31bd714f3af02b92f8991322fd42d9de9fcdc975 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 12 Aug 2026 20:11:43 -0500 Subject: [PATCH 3/4] Add a Metal kernel for the MLX incomplete gamma --- pytensor/link/mlx/dispatch/scalar/gammainc.py | 239 +++++++++++++++++- tests/link/mlx/scalar/test_gammainc.py | 85 ++++++- 2 files changed, 309 insertions(+), 15 deletions(-) diff --git a/pytensor/link/mlx/dispatch/scalar/gammainc.py b/pytensor/link/mlx/dispatch/scalar/gammainc.py index 7c40f92cf4..4b3f3aee06 100644 --- a/pytensor/link/mlx/dispatch/scalar/gammainc.py +++ b/pytensor/link/mlx/dispatch/scalar/gammainc.py @@ -1,9 +1,16 @@ +from collections.abc import Callable + import mlx.core as mx import numpy as np from pytensor.link.mlx.dispatch.basic import mlx_funcify +from pytensor.link.mlx.dispatch.scalar.erf import _METAL_HEADER, _metal_constants from pytensor.link.mlx.dispatch.scalar.helpers import _exp, _working_precision -from pytensor.link.mlx.dispatch.scalar.math import _lanczos_log_gamma +from pytensor.link.mlx.dispatch.scalar.math import ( + _LANCZOS_COEFFS, + _LANCZOS_G, + _lanczos_log_gamma, +) from pytensor.scalar.math import Erfc, GammaInc, GammaIncC @@ -352,25 +359,229 @@ def _incomplete_gamma(a, x, const, erfc, *, lower): return mx.where(saturated, const(1.0 if lower else 0.0), selected) -@mlx_funcify.register(GammaIncC) -def mlx_funcify_GammaIncC(op, **kwargs): - erfc = mlx_funcify(Erfc()) +# The vectorized form runs all three expansions for every element and discards two: +# timed separately under mx.compile at n = 1e6 they cost 1.6, 2.4 and 6.4 ms against +# 11.6 ms for the composite, so it is very nearly their sum. A kernel takes one branch +# and stops each series when it has converged, which for the ascending series is a +# median of 8 terms out of the 50 the graph must always unroll. +_METAL_GAMMAINC_HEADER = ( + _METAL_HEADER + + _metal_constants( + GI_LANCZOS=_LANCZOS_COEFFS, + GI_TEMME_D=tuple(value for row in _TEMME_D for value in row), + # The eta series coefficients are loop-invariant, so they are tabulated for the + # same reason the vectorized path folds them at trace time: written inline they + # are a division per term that only constant folding would remove + GI_ETA_C=tuple(2.0 / (k + 2.0) for k in range(_ETA_SERIES_TERMS)), + ) + + f""" +#define GI_SERIES_TERMS {_SERIES_TERMS} +#define GI_CF_TERMS {_CF_TERMS} +#define GI_ETA_TERMS {_ETA_SERIES_TERMS} +#define GI_ETA_MAX_SIGMA {_ETA_SERIES_MAX_SIGMA}f +#define GI_TEMME_MIN_A {_TEMME_MIN_A}f +#define GI_TEMME_MIN_R {_TEMME_MIN_RATIO}f +#define GI_TEMME_MAX_R {_TEMME_MAX_RATIO}f +#define GI_TEMME_ROWS {len(_TEMME_D)} +#define GI_TEMME_COLS {len(_TEMME_D[0])} +#define GI_LANCZOS_N {len(_LANCZOS_COEFFS)} +#define GI_LANCZOS_T {_LANCZOS_G + 0.5}f +#define GI_LOG_SQRT_2PI {float(0.5 * np.log(2.0 * np.pi))!r}f +#define GI_TINY {_TINY}f +// just under one float32 epsilon: below this a term no longer moves the sum +#define GI_CONVERGED 1.0e-7f + +static inline float gi_erfc(float value) {{ + float y = fabs(value); + float res = (y <= ERF_THRESH) ? 1.0f - erf_series(y) + : exp(-y * y) * erfcx_upper(y); + return (value < 0.0f) ? 2.0f - res : res; +}} + +// log Gamma(w + 1), shifted the same way the vectorized Lanczos helper is +static inline float gi_lgamma1p(float w) {{ + float series = GI_LANCZOS[0]; + for (int i = 1; i < GI_LANCZOS_N; ++i) series += GI_LANCZOS[i] / (w + (float)i); + float t = w + GI_LANCZOS_T; + return GI_LOG_SQRT_2PI + (w + 0.5f) * log(t) - t + log(series); +}} + +static inline float gi_series(float a, float x) {{ + // x - x rather than 0, so a nan argument still propagates + if (!(x > 0.0f)) return x - x; + float term = 1.0f; + float total = 1.0f; + for (int n = 1; n < GI_SERIES_TERMS; ++n) {{ + term *= x / (a + (float)n); + total += term; + if (fabs(term) <= GI_CONVERGED * fabs(total)) break; + }} + return total * exp(-x + a * log(x) - gi_lgamma1p(a)); +}} + +static inline float gi_cf(float a, float x) {{ + float b = x + 1.0f - a; + float c = 1.0f / GI_TINY; + float d = 1.0f / ((b == 0.0f) ? GI_TINY : b); + float h = d; + for (int i = 1; i < GI_CF_TERMS; ++i) {{ + float an = -(float)i * ((float)i - a); + b += 2.0f; + d = an * d + b; + if (fabs(d) < GI_TINY) d = GI_TINY; + c = b + an / c; + if (fabs(c) < GI_TINY) c = GI_TINY; + d = 1.0f / d; + float delta = d * c; + h *= delta; + if (fabs(delta - 1.0f) <= GI_CONVERGED) break; + }} + return h * exp(-x + a * log(x) - gi_lgamma1p(a - 1.0f)); +}} + +static inline float gi_eta(float a, float x) {{ + float sigma = (x - a) / a; + float value; + if (fabs(sigma) <= GI_ETA_MAX_SIGMA) {{ + float series = 0.0f; + for (int k = GI_ETA_TERMS - 1; k >= 0; --k) {{ + series = series * (-sigma) + GI_ETA_C[k]; + }} + value = sigma * sigma * series; + }} else {{ + value = -2.0f * (log1p(sigma) - sigma); + }} + return sign(sigma) * sqrt(fmax(value, 0.0f)); +}} + +static inline float gi_temme(float a, float x, float sgn) {{ + float eta = gi_eta(a, x); + float total = 0.0f; + float a_power = 1.0f; + float inverse_a = 1.0f / a; + for (int k = 0; k < GI_TEMME_ROWS; ++k) {{ + int row = k * GI_TEMME_COLS; + float coefficient = GI_TEMME_D[row + GI_TEMME_COLS - 1]; + for (int n = GI_TEMME_COLS - 2; n >= 0; --n) {{ + coefficient = coefficient * eta + GI_TEMME_D[row + n]; + }} + total += coefficient * a_power; + a_power *= inverse_a; + }} + float correction = exp(-0.5f * a * eta * eta) * rsqrt(2.0f * M_PI_F * a) * total; + return 0.5f * gi_erfc(sgn * eta * sqrt(a * 0.5f)) + sgn * correction; +}} +""" +) - def gammaincc(a, x): - z, const, out_dtype = _working_precision(x) - shape = mx.array(a).astype(z.dtype) - return _incomplete_gamma(shape, z, const, erfc, lower=False).astype(out_dtype) - return gammaincc +_METAL_GAMMAINC_SOURCE = """ + uint i = thread_position_in_grid.x; + float a = (float)a_buf[i]; + float x = (float)x_buf[i]; + float res; + if (isinf(x) && x > 0.0f) { + // the continued fraction would form (1 / inf) * inf here and return nan + res = GI_LOWER ? 1.0f : 0.0f; + } else if (a >= GI_TEMME_MIN_A && x >= GI_TEMME_MIN_R * a && x <= GI_TEMME_MAX_R * a) { + res = gi_temme(a, x, GI_SIGN); + } else if (x <= a + 1.0f) { + float lower = gi_series(a, x); + res = GI_LOWER ? lower : 1.0f - lower; + } else { + float upper = gi_cf(a, x); + res = GI_LOWER ? 1.0f - upper : upper; + } + out[i] = (T)res; +""" + + +_METAL_GAMMAINC_KERNELS: dict[bool, Callable | None] = {} + + +def _metal_gammainc_kernel(lower): + """Build and cache the Metal kernel for one tail, or None if it cannot be built. + + The guard catches a failure to construct the kernel object. It does not catch a + malformed source: Metal compiles lazily on first call and aborts the process rather + than raising, so the source has no runtime safety net and the fallback cannot stand + in for one. + """ + if lower not in _METAL_GAMMAINC_KERNELS: + # the trailing newline matters: MLX appends its own template declaration + # straight after this header, and a #define running into it fails to compile + header = ( + _METAL_GAMMAINC_HEADER + + f"#define GI_LOWER {int(lower)}\n" + + f"#define GI_SIGN {-1.0 if lower else 1.0}f\n" + ) + try: + _METAL_GAMMAINC_KERNELS[lower] = mx.fast.metal_kernel( + name=f"pytensor_gamma{'inc' if lower else 'incc'}", + input_names=["a_buf", "x_buf"], + output_names=["out"], + header=header, + source=_METAL_GAMMAINC_SOURCE, + ) + except Exception: + _METAL_GAMMAINC_KERNELS[lower] = None + return _METAL_GAMMAINC_KERNELS[lower] + + +def _metal_gammainc_call(a, x, lower): + """Evaluate one tail through its Metal kernel, or return None to fall back. + + Metal has no float64 and the kernel needs the GPU stream, which between them leave + exactly the case the vectorized path handles least well and float64 needs least: the + default device at single precision. + """ + if ( + a.dtype != mx.float32 + or x.dtype != mx.float32 + or not mx.metal.is_available() + or mx.default_device() != mx.gpu + ): + return None + kernel = _metal_gammainc_kernel(lower) + if kernel is None: + return None + # The kernel indexes two flat buffers in lockstep, so the operands are broadcast + # against each other first: the op admits a scalar shape parameter against an array + # of values, and a 0-d input binds as a value rather than a pointer and cannot be + # subscripted + shape, value = mx.broadcast_arrays(a, x) + flat_shape, flat_value = shape.reshape(-1), value.reshape(-1) + (out,) = kernel( + inputs=[flat_shape, flat_value], + template=[("T", flat_value.dtype)], + grid=(flat_value.size, 1, 1), + threadgroup=(min(256, max(flat_value.size, 1)), 1, 1), + output_shapes=[flat_value.shape], + output_dtypes=[flat_value.dtype], + ) + return out.reshape(shape.shape) -@mlx_funcify.register(GammaInc) -def mlx_funcify_GammaInc(op, **kwargs): +def _incomplete_gamma_dispatch(lower): + """Build the callable one of the two ops dispatches to.""" erfc = mlx_funcify(Erfc()) - def gammainc(a, x): + def incomplete_gamma(a, x): z, const, out_dtype = _working_precision(x) shape = mx.array(a).astype(z.dtype) - return _incomplete_gamma(shape, z, const, erfc, lower=True).astype(out_dtype) + out = _metal_gammainc_call(shape, z, lower) + if out is None: + out = _incomplete_gamma(shape, z, const, erfc, lower=lower) + return out.astype(out_dtype) - return gammainc + return incomplete_gamma + + +@mlx_funcify.register(GammaIncC) +def mlx_funcify_GammaIncC(op, **kwargs): + return _incomplete_gamma_dispatch(lower=False) + + +@mlx_funcify.register(GammaInc) +def mlx_funcify_GammaInc(op, **kwargs): + return _incomplete_gamma_dispatch(lower=True) diff --git a/tests/link/mlx/scalar/test_gammainc.py b/tests/link/mlx/scalar/test_gammainc.py index e1a71faed3..d71dd12d62 100644 --- a/tests/link/mlx/scalar/test_gammainc.py +++ b/tests/link/mlx/scalar/test_gammainc.py @@ -1,3 +1,4 @@ +import os from functools import partial import numpy as np @@ -5,7 +6,7 @@ import scipy.special import pytensor.tensor as pt -from pytensor.scalar.math import GammaInc, GammaIncC +from pytensor.scalar.math import Erfc, GammaInc, GammaIncC from pytensor.tensor.type import vector from tests.link.mlx.test_basic import compare_mlx_and_py @@ -17,7 +18,10 @@ _TEMME_MIN_A, _TEMME_MIN_RATIO, _TINY, + _incomplete_gamma, + _metal_gammainc_call, ) +from pytensor.link.mlx.dispatch.scalar.helpers import _working_precision def _sample(rng, shape_range, ratio_range, n): @@ -152,6 +156,85 @@ def test_incomplete_gamma_branch_boundaries(): ) +@pytest.mark.skipif( + not mlx.metal.is_available() or os.environ.get("PYTENSOR_MLX_SKIP_GPU") == "1", + reason="needs a GPU that can run kernels; set PYTENSOR_MLX_SKIP_GPU=1 where it cannot", +) +@pytest.mark.parametrize("lower", [True, False], ids=["P", "Q"]) +def test_incomplete_gamma_paths_agree(lower): + # Both tails are implemented twice: a Metal kernel, taken for float32 on the GPU + # stream, and the vectorized fallback taken everywhere else -- which is every float64 + # graph, since Metal has no float64 at all. Only the fallback runs on CI, so this is + # the only thing standing between the kernel and a silent drift. + # + # The kernel is not a transcription of the fallback: it takes one branch per element + # where the fallback evaluates all three, and it stops each series on convergence + # where the fallback always runs the full trip count. The sample has to reach every + # branch for that to mean anything, so it spans the ratio rather than the value. + rng = np.random.default_rng(71) + a_test_value, x_test_value = _sample(rng, (1e-2, 1e4), (1e-2, 20.0), 5001) + a_test_value = a_test_value.astype("float32") + x_test_value = x_test_value.astype("float32") + + with mlx.stream(mlx.gpu): + a_mlx = mlx.array(a_test_value) + x_mlx = mlx.array(x_test_value) + # Without this the test compares the fallback against itself and passes whatever + # the kernel does + kernel = _metal_gammainc_call(a_mlx, x_mlx, lower) + assert kernel is not None, ( + "the Metal kernel did not engage, so this comparison proves nothing" + ) + metal = np.asarray(kernel) + + with mlx.stream(mlx.cpu): + z, const, _ = _working_precision(mlx.array(x_test_value)) + vectorized = np.asarray( + _incomplete_gamma( + mlx.array(a_test_value), z, const, mlx_funcify(Erfc()), lower=lower + ) + ) + + # Both are float32, so agreement is claimed at float32 terms; below the smallest + # normal neither carries a relative answer to compare + np.testing.assert_allclose( + metal, vectorized, rtol=1e-3, atol=np.finfo("float32").tiny + ) + + +@pytest.mark.skipif( + not mlx.metal.is_available() or os.environ.get("PYTENSOR_MLX_SKIP_GPU") == "1", + reason="needs a GPU that can run kernels; set PYTENSOR_MLX_SKIP_GPU=1 where it cannot", +) +@pytest.mark.parametrize("lower", [True, False], ids=["P", "Q"]) +def test_incomplete_gamma_kernel_edge_cases(lower): + # The kernel handles every edge a second time, and nothing else reaches it with these + # arguments: test_incomplete_gamma_edge_cases runs on the default device, which the + # conftest pins to the CPU, and test_incomplete_gamma_paths_agree samples log-uniform + # ratios that never produce a zero, an infinity or a nan. That gap is how the kernel + # and the fallback both returned nan at x = inf. + a_test_value = np.array( + [1.0, 0.5, 1e-3, 200.0, 1.0, 200.0, 1.0, 0.5, 200.0, 1.0], dtype="float32" + ) + x_test_value = np.array( + [0.0, 0.0, 0.0, 0.0, 1e30, 1e30, np.inf, np.inf, np.inf, np.nan], + dtype="float32", + ) + + with mlx.stream(mlx.gpu): + kernel = _metal_gammainc_call( + mlx.array(a_test_value), mlx.array(x_test_value), lower + ) + assert kernel is not None, ( + "the Metal kernel did not engage, so this comparison proves nothing" + ) + metal = np.asarray(kernel) + + reference = scipy.special.gammainc if lower else scipy.special.gammaincc + truth = reference(a_test_value.astype("float64"), x_test_value.astype("float64")) + np.testing.assert_allclose(metal, truth, rtol=1e-5) + + def test_continued_fraction_floor_survives_float32(): # The continued fraction seeds one Lentz variable at 1 / _TINY and guards both against # it. A floor below float32's smallest normal -- 1e-300 was -- rounds to zero there, From 18811f726c0c9f83a635130a1e074125d63c6a94 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 12 Aug 2026 20:30:17 -0500 Subject: [PATCH 4/4] Cover the MLX incomplete gamma kernel dispatch routing --- tests/link/mlx/scalar/test_gammainc.py | 64 +++++++++++++++++++++++++- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/tests/link/mlx/scalar/test_gammainc.py b/tests/link/mlx/scalar/test_gammainc.py index d71dd12d62..17486be34d 100644 --- a/tests/link/mlx/scalar/test_gammainc.py +++ b/tests/link/mlx/scalar/test_gammainc.py @@ -173,6 +173,22 @@ def test_incomplete_gamma_paths_agree(lower): # branch for that to mean anything, so it spans the ratio rather than the value. rng = np.random.default_rng(71) a_test_value, x_test_value = _sample(rng, (1e-2, 1e4), (1e-2, 20.0), 5001) + + # Enforce the claim above rather than trusting it: a later edit to the ranges could + # empty a branch, and the test would keep passing while covering less than it says + in_temme = ( + (a_test_value >= _TEMME_MIN_A) + & (x_test_value >= _TEMME_MIN_RATIO * a_test_value) + & (x_test_value <= _TEMME_MAX_RATIO * a_test_value) + ) + in_series = ~in_temme & (x_test_value <= a_test_value + 1.0) + for name, reached in ( + ("temme", in_temme), + ("series", in_series), + ("continued fraction", ~in_temme & ~in_series), + ): + assert reached.mean() > 0.05, f"the sample barely reaches the {name} branch" + a_test_value = a_test_value.astype("float32") x_test_value = x_test_value.astype("float32") @@ -196,10 +212,54 @@ def test_incomplete_gamma_paths_agree(lower): ) # Both are float32, so agreement is claimed at float32 terms; below the smallest - # normal neither carries a relative answer to compare + # normal neither carries a relative answer to compare. The two disagree by 1.2e-4 at + # worst and by nothing at all at the median, so the bound sits just above the observed + # spread rather than at a round number that would let real drift through np.testing.assert_allclose( - metal, vectorized, rtol=1e-3, atol=np.finfo("float32").tiny + metal, vectorized, rtol=3e-4, atol=np.finfo("float32").tiny + ) + + +@pytest.mark.skipif( + not mlx.metal.is_available() or os.environ.get("PYTENSOR_MLX_SKIP_GPU") == "1", + reason="needs a GPU that can run kernels; set PYTENSOR_MLX_SKIP_GPU=1 where it cannot", +) +@pytest.mark.parametrize( + "op_class, lower", [(GammaInc, True), (GammaIncC, False)], ids=["P", "Q"] +) +def test_incomplete_gamma_dispatch_reaches_kernel(op_class, lower): + # Everything else that exercises the kernel calls it directly, and everything that + # goes through the registered dispatch runs on the CPU stream the conftest pins. So + # removing the kernel from the dispatch would leave the whole suite green while every + # float32 GPU graph silently took the path that is an order of magnitude slower. + a_test_value, x_test_value = _sample( + np.random.default_rng(5), (0.1, 1e3), (0.05, 10.0), 1001 + ) + a_test_value = a_test_value.astype("float32") + x_test_value = x_test_value.astype("float32") + + with mlx.stream(mlx.gpu): + a_mlx = mlx.array(a_test_value) + x_mlx = mlx.array(x_test_value) + dispatched = np.asarray(mlx_funcify(op_class())(a_mlx, x_mlx)) + kernel = _metal_gammainc_call(a_mlx, x_mlx, lower) + assert kernel is not None, "the Metal kernel refused this input" + kernel = np.asarray(kernel) + + with mlx.stream(mlx.cpu): + z, const, _ = _working_precision(mlx.array(x_test_value)) + vectorized = np.asarray( + _incomplete_gamma( + mlx.array(a_test_value), z, const, mlx_funcify(Erfc()), lower=lower + ) + ) + + # Bit equality is only a discriminator while the two paths actually differ; if they + # ever agree exactly this test would pass without checking anything + assert not np.array_equal(kernel, vectorized), ( + "the two paths agree bitwise, so this test cannot tell them apart" ) + np.testing.assert_array_equal(dispatched, kernel) @pytest.mark.skipif(