Add backend-agnostic lgamma op - #23541
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds the lgamma mathematical operation to Keras ops, implementing a backend-agnostic fallback using the Lanczos approximation, along with corresponding unit tests. The feedback highlights a few critical improvements: first, scipy.special must be imported locally in the new test cases to prevent NameError failures; second, native lgamma implementations should be added to the respective backend wrappers (TensorFlow, JAX, PyTorch, and NumPy) to prefer optimized native execution over the fallback; and finally, several mathematical constants should be precomputed rather than calculated at runtime to reduce symbolic overhead.
| def _lgamma(x): | ||
| if not config._use_backend_agnostic_ops() and hasattr( | ||
| backend.math, "lgamma" | ||
| ): | ||
| return backend.math.lgamma(x) |
There was a problem hiding this comment.
While the backend-agnostic Lanczos approximation is a great fallback, relying on it for all backends is inefficient and potentially less numerically stable than using native implementations. TensorFlow (tf.math.lgamma), JAX (jax.lax.lgamma), and PyTorch (torch.lgamma) all provide highly optimized, native lgamma implementations.
To align with existing repository patterns (such as erf, erfc, etc.), please implement lgamma in the respective backend wrappers (e.g., keras/src/backend/tensorflow/math.py, keras/src/backend/jax/math.py, keras/src/backend/torch/math.py, and keras/src/backend/numpy/math.py using scipy.special.gammaln) so that hasattr(backend.math, "lgamma") evaluates to True and the native implementations are preferred.
| _LOG_LANCZOS_GAMMA_PLUS_HALF = ops.cast( | ||
| ops.log(_LANCZOS_GAMMA_PLUS_HALF), compute_dtype | ||
| ) |
There was a problem hiding this comment.
Instead of computing ops.log on a constant value at runtime, we can precompute this constant to avoid unnecessary symbolic operations and runtime overhead.
_LOG_LANCZOS_GAMMA_PLUS_HALF = ops.cast(2.0149030205422647, compute_dtype)There was a problem hiding this comment.
I kept it this way to improve readability
| log_sqrt_two_pi = ops.cast( | ||
| (ops.log(2.0) + ops.log(_PI)) / 2.0, compute_dtype | ||
| ) |
There was a problem hiding this comment.
Instead of computing ops.log on constant values at runtime, we can precompute log(sqrt(2π)) as a constant to avoid unnecessary symbolic operations and runtime overhead.
log_sqrt_two_pi = ops.cast(0.9189385332046727, compute_dtype)| reflection = ops.where( | ||
| ops.isfinite(reflection_denom), | ||
| ops.cast(ops.log(_PI), compute_dtype) - reflection_denom - log_y, | ||
| -reflection_denom, | ||
| ) |
There was a problem hiding this comment.
Instead of computing ops.log(_PI) at runtime, we can precompute log(π) as a constant to avoid unnecessary symbolic operations and runtime overhead.
| reflection = ops.where( | |
| ops.isfinite(reflection_denom), | |
| ops.cast(ops.log(_PI), compute_dtype) - reflection_denom - log_y, | |
| -reflection_denom, | |
| ) | |
| reflection = ops.where( | |
| ops.isfinite(reflection_denom), | |
| ops.cast(1.1447298858494002, compute_dtype) - reflection_denom - log_y, | |
| -reflection_denom, | |
| ) |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23541 +/- ##
=======================================
Coverage 85.02% 85.03%
=======================================
Files 468 468
Lines 71087 71142 +55
Branches 11788 11792 +4
=======================================
+ Hits 60442 60494 +52
- Misses 7632 7634 +2
- Partials 3013 3014 +1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
hertschuh
left a comment
hertschuh
left a comment
There was a problem hiding this comment.
Thanks for adding this!
Can you also add the backend specific implementations?
- Numpy:
scipy.special.gammaln - JAX:
jax.scipy.special.gammaln - Tensorflow:
tf.math.lgamma - Torch:
torch.lgamma
I think they'll be faster. But the fallback will still be used by OpenVino and new backends.
|
|
||
| # If the input is less than 0.5 use Euler's reflection formula: | ||
| # gamma(x) = pi / (sin(pi * x) * gamma(1 - x)) | ||
| need_to_reflect = x < 0.5 |
There was a problem hiding this comment.
You're using a lot of Python built-in ops (+, /, - etc.). I think those are mostly fine (i.e. they work with native tensors in all backends).
But for the comparison, I'd rather use the op:
need_to_reflect = ops.less(x, 0.5)| _LANCZOS_GAMMA_PLUS_HALF = _LANCZOS_GAMMA + 0.5 | ||
| _LOG_LANCZOS_GAMMA_PLUS_HALF = ops.cast( | ||
| ops.log(_LANCZOS_GAMMA_PLUS_HALF), compute_dtype | ||
| ) | ||
|
|
||
| t = z + _LANCZOS_GAMMA_PLUS_HALF | ||
| log_t = _LOG_LANCZOS_GAMMA_PLUS_HALF + ops.log1p( | ||
| z / _LANCZOS_GAMMA_PLUS_HALF | ||
| ) |
There was a problem hiding this comment.
Up until here, the constants were consistently cast to compute_dtype. But starting from here, some casts are missing.
For instance when you do t = z + _LANCZOS_GAMMA_PLUS_HALF, the _LANCZOS_GAMMA_PLUS_HALF is still the raw Python float64 and wasn't cast to compute_dtype.
Same line 1465 z / _LANCZOS_GAMMA_PLUS_HALF.
There was a problem hiding this comment.
Thanks for catching that! Done!
| # log(sqrt(2π)) = (log(2) + log(pi)) / 2 | ||
| log_sqrt_two_pi = ops.cast( | ||
| (ops.log(2.0) + ops.log(_PI)) / 2.0, compute_dtype | ||
| ) |
There was a problem hiding this comment.
I agree with you on the readability of not hardcoding a constant for this. The flip side is that we don't want to re-compute the constants on each call. So here is a compromise. It computes the constants once for all using the Python math function (import math).
line 1433:
_LOG_SQRT_TWO_PI = (math.log(2.0) + math.log(_PI)) / 2.0Then
log_sqrt_two_pi = ops.cast(_LOG_SQRT_TWO_PI, compute_dtype)You can do the same with _LANCZOS_GAMMA_PLUS_HALF and _LOG_LANCZOS_GAMMA_PLUS_HALF.
|
|
||
| abs_x = ops.abs(x) | ||
| abs_frac_x = abs_x - ops.floor(abs_x) | ||
| reduced_frac_x = ops.where(abs_frac_x > 0.5, 1.0 - abs_frac_x, abs_frac_x) |
There was a problem hiding this comment.
ops.greater(abs_frac_x, 0.5)
SamanehSaadat
left a comment
SamanehSaadat
left a comment
There was a problem hiding this comment.
Thanks for the review, Fabien!
Is it okay if I add backend-specific implementations in a separate PR and keep this one focused on the backend-agnostic implementation?
|
|
||
| # If the input is less than 0.5 use Euler's reflection formula: | ||
| # gamma(x) = pi / (sin(pi * x) * gamma(1 - x)) | ||
| need_to_reflect = x < 0.5 |
| _LANCZOS_GAMMA_PLUS_HALF = _LANCZOS_GAMMA + 0.5 | ||
| _LOG_LANCZOS_GAMMA_PLUS_HALF = ops.cast( | ||
| ops.log(_LANCZOS_GAMMA_PLUS_HALF), compute_dtype | ||
| ) | ||
|
|
||
| t = z + _LANCZOS_GAMMA_PLUS_HALF | ||
| log_t = _LOG_LANCZOS_GAMMA_PLUS_HALF + ops.log1p( | ||
| z / _LANCZOS_GAMMA_PLUS_HALF | ||
| ) |
There was a problem hiding this comment.
Thanks for catching that! Done!
| # log(sqrt(2π)) = (log(2) + log(pi)) / 2 | ||
| log_sqrt_two_pi = ops.cast( | ||
| (ops.log(2.0) + ops.log(_PI)) / 2.0, compute_dtype | ||
| ) |
|
|
||
| abs_x = ops.abs(x) | ||
| abs_frac_x = abs_x - ops.floor(abs_x) | ||
| reduced_frac_x = ops.where(abs_frac_x > 0.5, 1.0 - abs_frac_x, abs_frac_x) |
google-ml-butler
Bot
added
kokoro:force-run
ready to pull
4 participants
Description
This PR adds backend-agnostic
lgammaop. Note that thelgammaop doesn't have backend-specific implementations.Contributor Agreement
Please review our AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:
Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.