Reduce redundant digamma and lgamma evaluations in special functions. #41869

Merged
copybara-service[bot] opened 4:53pm on October 21, 2026 wanted to merge 122 Ξ” into google/jax main from
test_1004201622
copybara-service[bot] authored
of work during October 22
Diff Delta:
122
About 9 Diff Delta/hour
Classified as:  General

copybara-service[bot]'s Description of Work

Reduce redundant digamma and lgamma evaluations in special functions.

igamma_grad_a_impl and random_gamma_grad_impl evaluated both digamma(a) and digamma(a + 1) across unconditionally evaluated select branches even though each element only uses one branch after the loop finishes. Select a or a + 1 up front and pass the shared log(x) - digamma(...) term into both helpers, cutting digamma calls from two to one in each function.

In _hyp2f1_digamma_transform, AMS55 15.3.10–15.3.12 only apply when d = c - a - b is within 50 * eps of an integer rd, while a 0..|rd| loop already runs to accumulate log(|rd|!). Accumulate the harmonic sum 1 / (e - t) inside that loop with a first-order trigamma correction for e - |rd|, replacing digamma(1 + e) and cutting digamma calls from three to two. Also hoist the duplicate lax.lgamma(new_a) call in betaln.

3 total changed files
Loading changes...
Loading changes...
Loading changes...
Review complete πŸ’ͺ
No additional context awaits in the abyss below