FlatPPL alpha 15

Alpha wheels for CPython 3.12, 3.13, 3.14 and 3.15 on macOS ARM64, Linux x86-64 and Linux ARM64. JAX and jaxlib are pinned to 0.12.0.1.

• Accurate float32 lgamma and digamma (#42). A new pass, lower-chlo-gamma-f32, rewrites float32 chlo.lgamma with the JAX 0.12 float32 algorithm, bit for bit, and float32 chlo.digamma with a root-anchored minimax rewrite. hlo_call runs the pass after AD on every path, so a value and its gradient use the same lgamma. float64 is unchanged.
• Against mpmath on 100k float32 points (M4 CPU), the maximum error falls:
  • lgamma: from 3.9e6 ulp near x = 2 (and unbounded at x = 1) to 2.48 ulp.
  • digamma: from 335306 ulp near its root to 4.08 ulp.
  • GH200 and GB10 agree with M4 within one ulp, for example lgamma 2.48 ulp and digamma 4.08 ulp at most on positive arguments.
• Near the negative roots of digamma the error stays large. Its mean and 99th percentile are lower than before, but on GPU the single worst point in some unit intervals is higher.
• Speed: one float32 lgamma value costs about 1.6 times as much as before, which is what jax.lax.lgamma costs since JAX 0.12. digamma and lgamma gradients run about 25% faster. On five posteriordb models the gradient time does not change measurably.
• Otherwise identical to alpha14: upstream EnzymeAD/Enzyme-JAX main through 2cc7c1b, manylinux_2_28 on both Linux architectures, macOS 14.

Qualification with JAX 0.12.0.1, every ABI on every platform:

• The ABI check passes.
• The x64 feature tests pass: 72 passed and 8 skipped on macOS, 71 passed and 9 skipped on Linux ARM64, and 70 passed and 9 skipped on Linux x86-64. The x86-64 run also fails the C++ cpp_call test, as with alpha14, because that host's system C++ headers are not discoverable (see below).
• hlo_clamping and hlo_linearize pass: 32 tests and 8 subtests.
• On Linux x86-64, the C++ checks need the host's system C++ headers to be discoverable, as with alpha14.
• The FlatPPL StableHLO suite gives the same results as alpha14, on one flatppl-rust build (main 109efb3): 1100 passed and 1 skipped on macOS, and 1100 passed and 1 skipped on Linux x86-64 for both alpha14 and alpha15.

Each wheel differs from its qualified staging build only in METADATA and RECORD. SHA256SUMS and provenance.json identify the wheel payloads and the source commit.