Improve accuracy of lax.digamma to within 1 ULP across all float types. #41867

Merged
copybara-service[bot] opened 4:43pm on October 21, 2026 wanted to merge 1078 Ξ” into google/jax main from
test_1000339049

Pull Request Overview

  • Opened on October 21, 2026
  • Status Merged
  • Commit count 1 with first commit October 22, 2026

Total Delta

1078 Total Diff Delta

Open Days

Open 17 weekdays

Test Delta

118 Diff Delta in Test Files
Breakdown by Phase

How long has this pull request spent in each phase of its lifecycle?

Data pending calculation for pull request

Author avatar

Improve accuracy of lax.digamma to within 1 ULP across all float types.

Improve accuracy of lax.digamma to within 1 ULP across all float types.

The previous CHLO lowering of lax.digamma (jax.scipy.special.digamma) loses accuracy near zeros of digamma where subtraction cancels leading digits, as well as on accelerators that flush subnormal numbers to zero or lack fused multiply-add.

Replace the CHLO lowering of lax.digamma for bfloat16, float16, float32, and float64 with digamma_impl in jax._src.lax.digamma, keeping error within 1.0 ULP for float64, float32, and float16 and within 0.5 ULP for bfloat16 on CPU, GPU, and TPU. Both float32 and float64 shift small arguments on [0.5, z_min) upward with the digamma recurrence, evaluate the asymptotic Bernoulli series on larger arguments, and reflect arguments below 0.5 with Euler's reflection formula after subtracting the nearest cotangent poles at 0, Β±1, and Β±2.

On CPU, float32 uses native float64 intermediates. On GPU and TPU in float32, and on all backends in float64, intermediate values are tracked as (hi, lo) floating-point pairs without requiring hardware FMA or subnormals, supplemented by a local Taylor polynomial around the positive zero 1.461632145 and precomputed outputs at the handful of negative inputs where cancellation exceeds working precision (three inputs in float32 and ten in float64). Shared (hi, lo) arithmetic primitives and a four-interval log_f2 kernel live in jax._src.lax.double_word and are shared with jax._src.lax.lgamma.

PR was closed without comments.