Compute the Digamma function using Lanczos' approximation from "A Precision Approximation of the Gamma Function". SIAM Journal on Numerical Analysis series B. Vol. 1: digamma(z + 1) = log(t(z)) + A'(z) / A(z) - kLanczosGamma / t(z) t(z) = z + kLanczosGamma + 1/2 A(z) = kBaseLanczosCoeff + sigma(k = 1, n, kLanczosCoefficients[i] / (z + k)) A'(z) = sigma(k = 1, n, kLanczosCoefficients[i] / (z + k) /
| 626 | // A(z) = kBaseLanczosCoeff + sigma(k = 1, n, kLanczosCoefficients[i] / (z + k)) |
| 627 | // A'(z) = sigma(k = 1, n, kLanczosCoefficients[i] / (z + k) / (z + k)) |
| 628 | XlaOp Digamma(XlaOp input) { |
| 629 | auto do_it = [](XlaOp input) { |
| 630 | XlaOp zero = ScalarLike(input, 0); |
| 631 | XlaOp one_half = ScalarLike(input, 0.5); |
| 632 | XlaOp one = ScalarLike(input, 1); |
| 633 | |
| 634 | XlaOp pi = ScalarLike(input, M_PI); |
| 635 | |
| 636 | XlaOp lanczos_gamma = ScalarLike(input, kLanczosGamma); |
| 637 | XlaOp lanczos_gamma_plus_one_half = ScalarLike(input, kLanczosGamma + 0.5); |
| 638 | XlaOp log_lanczos_gamma_plus_one_half = |
| 639 | ScalarLike(input, std::log(kLanczosGamma + 0.5)); |
| 640 | |
| 641 | XlaOp base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); |
| 642 | |
| 643 | // If the input is less than 0.5 use Euler's reflection formula: |
| 644 | // digamma(x) = digamma(1 - x) - pi * cot(pi * x) |
| 645 | XlaOp need_to_reflect = Lt(input, one_half); |
| 646 | XlaOp z = Select(need_to_reflect, -input, input - one); |
| 647 | |
| 648 | XlaOp num = zero; |
| 649 | XlaOp denom = base_lanczos_coeff; |
| 650 | for (int i = 0; i < kLanczosCoefficients.size(); ++i) { |
| 651 | XlaOp lanczos_coefficient = ScalarLike(input, kLanczosCoefficients[i]); |
| 652 | XlaOp index = ScalarLike(input, i); |
| 653 | num = num - lanczos_coefficient / ((z + index + one) * (z + index + one)); |
| 654 | denom = denom + lanczos_coefficient / (z + index + one); |
| 655 | } |
| 656 | |
| 657 | // To improve accuracy on platforms with less-precise log implementations, |
| 658 | // compute log(lanczos_gamma_plus_one_half) at compile time and use log1p on |
| 659 | // the device. |
| 660 | // log(t) = log(kLanczosGamma + 0.5 + z) |
| 661 | // = log(kLanczosGamma + 0.5) + log1p(z / (kLanczosGamma + 0.5)) |
| 662 | XlaOp t = lanczos_gamma_plus_one_half + z; |
| 663 | XlaOp log_t = log_lanczos_gamma_plus_one_half + |
| 664 | Log1p(z / lanczos_gamma_plus_one_half); |
| 665 | |
| 666 | XlaOp y = log_t + num / denom - lanczos_gamma / t; |
| 667 | |
| 668 | // We need to be careful how we compute cot(pi * input) below: For |
| 669 | // near-integral values of `input`, pi * input can lose precision. |
| 670 | // |
| 671 | // Input is already known to be less than 0.5 (otherwise we don't have to |
| 672 | // reflect). We shift values smaller than -0.5 into the range [-.5, .5] to |
| 673 | // increase precision of pi * input and the resulting cotangent. |
| 674 | XlaOp reduced_input = input + Abs(Floor(input + ScalarLike(input, 0.5))); |
| 675 | XlaOp reflection = |
| 676 | y - pi * Cos(pi * reduced_input) / Sin(pi * reduced_input); |
| 677 | XlaOp real_result = Select(need_to_reflect, reflection, y); |
| 678 | |
| 679 | // Digamma has poles at negative integers and zero; return nan for those. |
| 680 | return Select(And(Le(input, zero), Eq(input, Floor(input))), |
| 681 | FullLike(input, std::numeric_limits<float>::quiet_NaN()), |
| 682 | real_result); |
| 683 | }; |
| 684 | |
| 685 | auto& b = *input.builder(); |