betaln(a, b) is too large by exactly min(a, b) when min(a, b)/max(a, b) is subnormal

Description

jax.scipy.special.betaln(a, b) is exact across most of its range, but is too large by exactly min(a, b) once min(a, b) / max(a, b) falls below the smallest normal double. scipy.special.betaln is correct throughout.

For betaln(2, N) the first wrong result is one ULP above N = 2**1023, which is exactly where 2 / N stops being a normal double:

N = 8.988465674311579e+307  (2**1023 - 1ulp)   a/b normal      jax - scipy = +0.00
N = 8.98846567431158e+307   (2**1023)          a/b normal      jax - scipy = +0.00
N = 8.988465674311582e+307  (2**1023 + 1ulp)   a/b subnormal   jax - scipy = +2.00

The downstream consequence is silent and large. 1 / ((N + 1) * B(N - k + 1, k + 1)) is the standard cancellation-free way to evaluate a binomial coefficient, and it comes out a factor of e**-2 — 86% — low for every N above 2**1023.

Reproducing code

import jax

jax.config.update("jax_enable_x64", True)
import numpy as np
import scipy.special as ss
from jax.scipy.special import betaln

tiny = np.finfo(np.float64).tiny
cases = [(1.0, 1e308), (2.0, 1e308), (1.5, 1e308), (2.0, 8.99e307),
         (2.0, 8.9e307), (3.0, 1e308), (2.0, 1e300)]

print(f"{'a':>5} {'b':>12} {'a/b':>12} {'subnormal':>10} {'jax':>16} {'scipy':>16} {'err':>7}")
for a, b in cases:
    h, j, s = a / b, float(betaln(a, b)), float(ss.betaln(a, b))
    print(f"{a:>5} {b:>12.3e} {h:>12.4e} {str(0 < h < tiny):>10} {j:>16.8f} {s:>16.8f} {j - s:>+7.2f}")

    a            b          a/b  subnormal              jax            scipy     err
  1.0   1.000e+308  1.0000e-308       True    -708.19620864    -709.19620864   +1.00
  2.0   1.000e+308  2.0000e-308       True   -1416.39241728   -1418.39241728   +2.00
  1.5   1.000e+308  1.5000e-308       True   -1062.41509520   -1063.91509520   +1.50
  2.0   8.990e+307  2.2247e-308       True   -1416.17947280   -1418.17947280   +2.00
  2.0   8.900e+307  2.2472e-308      False   -1418.15934965   -1418.15934965   +0.00
  3.0   1.000e+308  3.0000e-308      False   -2126.89547875   -2126.89547875   +0.00
  2.0   1.000e+300  2.0000e-300      False   -1381.55105580   -1381.55105580   +0.00

The error is exactly a whenever a / b is subnormal and exactly zero otherwise. Checked against mpmath at 400 digits rather than against SciPy alone — loggamma(1e308) is about 7e310, so resolving the ~1400 being measured needs more than 311 digits of working precision, and a lower dps silently reports zero error.

Cause

jax/_src/third_party/scipy/betaln.py, in algdiv:

h = a / b
...
d = b + (a - 0.5)
...
u = d * lax.log1p(a / b)

XLA flushes subnormals to zero, so once a / b is subnormal the quotient becomes 0, log1p(0) is 0, and u is dropped entirely — where its true value is d * (a/b) ≈ b * (a/b) = a. Hence an error of exactly a, and since betaln sorts its arguments with a, b = minimum(a, b), maximum(a, b), exactly min(a, b) as seen by the caller.

The flush is directly observable:

>>> float(1.0 / jnp.float64(1e308))
0.0
>>> 1.0 / np.float64(1e308)     # numpy, for contrast
1e-308

The Fortran algdiv this is derived from does not flush, which is why scipy.special.betaln is unaffected.

Suggested fix

Where h has flushed to zero, log1p(h) would have been h to the last bit anyway, so u = d * h = a * (d / b) — and d / b is close to 1, so it cannot itself underflow. That makes it one line, reached only when h == 0:

u = jnp.where(h == 0, a * (d / b), d * lax.log1p(h))

I ran this against mpmath at 400 digits:

• all flushed cases go to 0 ulp — (1, 1e308), (2, 1e308), (1.5, 1e308), (2, 8.99e307), (0.5, 1.7e308) are each bit-exact afterwards, and for (1.5, 1e308) that is one ulp better than SciPy;
• 500 random non-flushed pairs spanning a ∈ (1e-3, 50), b ∈ (8, 1e300) are bit-identical to the current output.

Note that the line also has to read h rather than recomputing a / b, which is what it does today.

Separate observation, from the same measurements

The b >= 8 branch is noticeably less accurate than SciPy's just above the cut, which is unrelated to the flush but showed up in the same sweep:

    a      b       jax err     scipy err
  2.0    7.9     -5.33e-15      8.88e-16     <- lgamma branch
  1.0    8.0      9.30e-07      0.00e+00     <- algdiv branch
  2.0    8.0      1.29e-06      0.00e+00
  3.0    8.0      1.34e-06      0.00e+00
  0.5    8.0      5.58e-07      0.00e+00
  5.0    9.0      7.73e-07      0.00e+00
  8.0   20.0      8.49e-08      0.00e+00
  2.0  100.0      1.05e-10      0.00e+00

Accuracy drops by about nine orders of magnitude crossing b = 8 into algdiv, then recovers as b grows. Happy to split this out into its own issue if it's worth tracking separately.

System info (python version, jaxlib version, accelerator, etc.)

jax:        0.11.1
jaxlib:     0.11.1
scipy:      1.18.1
python:     3.14.3
platform:   macOS-26.6.2-arm64
device:     CpuDevice(id=0), x64 enabled

Context: the algdiv path dates from the fix for #13267. Downstream report: https://github.com/JAXtronomy/spexial/issues/27