Add MLX dispatch and a Metal kernel for gammainc and gammaincc

[jessegrabowski]

gammainc and gammaincc are the last ops between MLX and logcdf for Gamma, InverseGamma, Poisson and ChiSquared.

The usual series-below-a+1, continued-fraction-above split does not port. Both need a term count growing like sqrt(a), and a here is a runtime tensor -- a shape parameter, usually a random variable -- so there is no count that is both fixed at graph-build time and sufficient at a = 50000. Temme's uniform asymptotic expansion (DLMF 8.12) covers the transition instead: its cost does not depend on a, and its accuracy improves as a grows, which is exactly where the other two fall apart. The d[k][n] table is generated by the script scipy ships at scipy/special/_precompute/gammainc_asy.py rather than transcribed.

Neither tail is ever taken as one minus the other in the direction that cancels. At a = 200, x = 0.5a the true P is 9.3e-19, so a complement reports zero with 100% error, and that point is nowhere near an obvious edge case. Temme supplies either tail directly by erfc(z) + erfc(-z) = 2, so flipping one sign moves it from Q to P with no subtraction anywhere; the series and the fraction each own the region where the tail they compute natively is the small one.

Accuracy against scipy over a in [1e-2, 1e5] and x/a in [1e-3, 50], ~30k points per tail:

            worst     median      99th
float64 P   2.2e-11   4.4e-16   8.5e-13
float64 Q   8.6e-12   0.0e+00   4.2e-13
float32 P   1.9e-04   9.9e-08   1.4e-05
float32 Q   1.9e-04   1.3e-08   2.5e-05

float32 is limited by the log prefactor, which subtracts two numbers of order a log a -- at a = 200, x = 2a those are both near 800 and differ by 60. Temme has no such subtraction, which is what sets its window at [0.4a, 1.8a]: widening it takes float32 from 7.9e-4 to 2.2e-4 and costs float64 8.1e-12 to 2.2e-11, and past that float64 degrades while float32 stops improving.

There is also a Metal kernel, because the vectorized form evaluates 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. The kernel takes one branch and stops each series on convergence, which for the ascending series is a median of 8 terms out of the 50 the graph must always unroll.

Both MLX columns are float32 on the GPU stream over identical buffers, and differ only in whether the dispatch may reach its kernel. scipy ships no float32 kernel for these, so that column is float64 and the comparison against it is not apples to apples on dtype -- passes is against one fused elementwise pass over the same buffers, which is the number that does not depend on scipy at all.

        n     kernel   vector    scipy   vs vector   passes
P   100,000     0.25     1.43     7.67        5.7x    2.3 vs  13
P 1,000,000     1.03    11.12    76.13       10.8x    4.8 vs  52
P10,000,000     9.66   137.92   769.82       14.3x   11.8 vs 168
Q   100,000     0.37     1.31     8.90        3.5x    3.4 vs  12
Q 1,000,000     1.04    10.20    90.68        9.8x    4.8 vs  48
Q10,000,000     9.61   142.12   906.59       14.8x   11.7 vs 173

Times are medians in ms from pytest-benchmark; IQR is 1-4% on every kernel row except n = 1e5, where both MLX columns are dispatch-bound and scatter by a third.

float64 has no kernel, since Metal has no float64, and runs about 2-3x slower than scipy rather than faster. That is structural: scipy iterates with an early exit where a branch-free graph runs every trip count for every element.

Two things worth a reviewer's attention. Metal cannot execute on CI at all, so the tests comparing the kernel against the vectorized path only ever run on a developer machine, and the vectorized implementation stays the reference rather than being retired in the kernel's favor. And a malformed kernel source aborts the process rather than raising -- Metal compiles lazily on first call -- so the try/except around building the kernel guards construction only and cannot stand in for a runtime fallback.

Part of #2095.

[ricardoV94]

There is also a Metal kernel, because the vectorized form evaluates 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. The kernel takes one branch and stops each series on convergence, which for the ascending series is a median of 8 terms out of the 50 the graph must always unroll.

What is a Meta kernel, and the time thing mentioned here is it measured on a heterogenous vector or a single scalar?

[jessegrabowski]

What is a Metal kernel

The C code strings