Compute the Lgamma function using Lanczos' approximation from "A Precision Approximation of the Gamma Function". SIAM Journal on Numerical Analysis series B. Vol. 1: lgamma(z + 1) = (log(2) + log(pi)) / 2 + (z + 1/2) * log(t(z)) - t(z) + A(z) t(z) = z + kLanczosGamma + 1/2 A(z) = kBaseLanczosCoeff + sigma(k = 1, n, kLanczosCoefficients[i] / (z + k))
| 490 | // t(z) = z + kLanczosGamma + 1/2 |
| 491 | // A(z) = kBaseLanczosCoeff + sigma(k = 1, n, kLanczosCoefficients[i] / (z + k)) |
| 492 | XlaOp Lgamma(XlaOp input) { |
| 493 | auto do_it = [](XlaOp input) { |
| 494 | XlaOp one_half = ScalarLike(input, 0.5); |
| 495 | XlaOp one = ScalarLike(input, 1); |
| 496 | |
| 497 | XlaOp pi = ScalarLike(input, M_PI); |
| 498 | XlaOp log_pi = ScalarLike(input, std::log(M_PI)); |
| 499 | XlaOp log_sqrt_two_pi = |
| 500 | ScalarLike(input, (std::log(2) + std::log(M_PI)) / 2); |
| 501 | |
| 502 | XlaOp lanczos_gamma_plus_one_half = ScalarLike(input, kLanczosGamma + 0.5); |
| 503 | XlaOp log_lanczos_gamma_plus_one_half = |
| 504 | ScalarLike(input, std::log(kLanczosGamma + 0.5)); |
| 505 | |
| 506 | XlaOp base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); |
| 507 | |
| 508 | // If the input is less than 0.5 use Euler's reflection formula: |
| 509 | // gamma(x) = pi / (sin(pi * x) * gamma(1 - x)) |
| 510 | XlaOp need_to_reflect = Lt(input, one_half); |
| 511 | XlaOp z = Select(need_to_reflect, -input, input - one); |
| 512 | |
| 513 | XlaOp x = base_lanczos_coeff; |
| 514 | for (int i = 0; i < kLanczosCoefficients.size(); ++i) { |
| 515 | XlaOp lanczos_coefficient = ScalarLike(input, kLanczosCoefficients[i]); |
| 516 | XlaOp index = ScalarLike(input, i); |
| 517 | x = x + lanczos_coefficient / (z + index + one); |
| 518 | } |
| 519 | |
| 520 | // To improve accuracy on platforms with less-precise log implementations, |
| 521 | // compute log(lanczos_gamma_plus_one_half) at compile time and use log1p on |
| 522 | // the device. |
| 523 | // log(t) = log(kLanczosGamma + 0.5 + z) |
| 524 | // = log(kLanczosGamma + 0.5) + log1p(z / (kLanczosGamma + 0.5)) |
| 525 | XlaOp t = lanczos_gamma_plus_one_half + z; |
| 526 | XlaOp log_t = log_lanczos_gamma_plus_one_half + |
| 527 | Log1p(z / lanczos_gamma_plus_one_half); |
| 528 | |
| 529 | // Compute the final result (modulo reflection). t(z) may be large, and we |
| 530 | // need to be careful not to overflow to infinity in the first term of |
| 531 | // |
| 532 | // (z + 1/2) * log(t(z)) - t(z). |
| 533 | // |
| 534 | // Therefore we compute this as |
| 535 | // |
| 536 | // (z + 1/2 - t(z) / log(t(z))) * log(t(z)). |
| 537 | // |
| 538 | XlaOp log_y = log_sqrt_two_pi + (z + one_half - t / log_t) * log_t + Log(x); |
| 539 | |
| 540 | // Compute the reflected value, used when x < 0.5: |
| 541 | // |
| 542 | // lgamma(x) = log(pi) - lgamma(1-x) - log(abs(sin(pi * x))). |
| 543 | // |
| 544 | // (The abs is because lgamma is the log of the absolute value of the gamma |
| 545 | // function.) |
| 546 | // |
| 547 | // We have to be careful when computing the final term above. gamma(x) goes |
| 548 | // to +/-inf at every integer x < 0, and this is controlled by the |
| 549 | // sin(pi * x) term. The slope is large, so precision is particularly |