Truncated gamma - #1187
Open
quattro wants to merge 8 commits into
Open
Truncated gamma#1187quattro wants to merge 8 commits into
quattro wants to merge 8 commits into
Conversation
fehiepsi
reviewed
Oct 14, 2021
|
|
||
| def icdf(self, q): | ||
| # https://github.com/pyro-ppl/numpyro/issues/969 | ||
| from numpyro.distributions.util import gammaincinv |
Member
There was a problem hiding this comment.
I think you can move this import to the top.
| return cls(batch_shape=aux_data) | ||
|
|
||
|
|
||
| def TruncatedGamma(base_gamma, low=None, high=None, validate_args=None): |
Member
There was a problem hiding this comment.
I think it is better to expose the parameters of Gamma here (TruncatedGamma(concentration, rate, low=..., high=...), rather than using a nested pattern. There are a couple of benefits with that:
- parameters of the distribution is defined probably in
args_constraints - it is easier to test
- no need to have flatten/unflatten logic
| base_gamma = Gamma.tree_unflatten(base_aux, base_flatten) | ||
| return cls(base_gamma, low=low) | ||
|
|
||
| @validate_sample |
Member
There was a problem hiding this comment.
Unfortunately, currently validate_sample logic does not work with cdf :(
| # until jax/lax has direct implementation we'll need to rely on tfp | ||
| # https://github.com/pyro-ppl/numpyro/issues/969 | ||
| try: | ||
| import tensorflow_probability as tfpm |
Member
There was a problem hiding this comment.
I think you can import tensorflow_probability.substrates.jax directly, to make sure that jax backend is installed.
| return lprob - jnp.log(1.0 - lscale) | ||
|
|
||
| def _scale_moment(self, t): | ||
| assert t > -self.base_gamma.concentration |
Member
There was a problem hiding this comment.
This won't work for jax arrays (which might have abstract values under jit compiling). You can use jnp.where to mask out the invalid cases like this.
| def log_prob(self, value): | ||
| lprob = self.base_gamma.log_prob(value) | ||
| lscale = self.base_gamma.cdf(self.low) | ||
| return lprob - jnp.log(1.0 - lscale) |
Member
There was a problem hiding this comment.
You can use log1p(-lscale) for a better numerical result
Member
|
@quattro Looking the the PR is is the good shape - just have small comments above. Any chance we can have this in the next numpyro release? |
Contributor
Author
|
Will try my best. Should have some time closer to Thanksgiving holidays, does that fall before next release schedule? |
Member
|
Absolutely, there is no plan for the release date yet. Thank you! |
|
Will we have this feature in the future? |
This was referenced Aug 23, 2026
Closed
deeb01
added a commit
to deeb01/numpyro
that referenced
this pull request
Sep 2, 2026
Add LeftTruncatedGamma, RightTruncatedGamma and TwoSidedTruncatedGamma, dispatched by TruncatedGamma(concentration, rate, low=, high=), following the structure outlined in pyro-ppl#969. Supersedes the inactive pyro-ppl#1187. The two-sided normalizer switches to the upper-tail form when the interval lies above the median: the lower-tail difference F(high) - F(low) is exactly zero in single precision for Gamma(2, 1) on [30, 40]. Adds gammainccinv to distributions/util.py, wrapping tfp.math.igammacinv.
juanitorduz
pushed a commit
that referenced
this pull request
Sep 12, 2026
* Add TruncatedGamma distributions Add LeftTruncatedGamma, RightTruncatedGamma and TwoSidedTruncatedGamma, dispatched by TruncatedGamma(concentration, rate, low=, high=), following the structure outlined in #969. Supersedes the inactive #1187. The two-sided normalizer switches to the upper-tail form when the interval lies above the median: the lower-tail difference F(high) - F(low) is exactly zero in single precision for Gamma(2, 1) on [30, 40]. Adds gammainccinv to distributions/util.py, wrapping tfp.math.igammacinv. * Address review: per-order tail switch, bisection icdf, guards - `_partial_moment` re-decides the tail switch at `concentration + order`. Gamma(alpha+k) has a larger median, so an interval above the base median can sit deep in the lower tail at the shifted order, where the upper-tail form is the one that cancels. Reusing the order-zero branch collapsed the second moment to zero and produced a negative variance. - `icdf` bisects the cdf instead of calling an inverse incomplete gamma, which saturates below a tail probability of ~3e-8 and returned draws outside the support once the retained mass got small. A Newton step from the bracketed root supplies the implicit derivative, so sampling stays differentiable. This removes the need for the `gammainccinv` wrapper, so `util.py` is untouched. - `mean` and `variance` now return nan for an underflowed normalizer, matching `cdf`; documented why that differs from `log_prob` returning -inf. - Documented the float32 cancellation in `E[X^2] - E[X]^2` for narrow intervals far from the origin. - Fixed two line-wrapped reST cross-references, gave the base-distribution assert a message pointing at TruncatedGamma, dropped a test assertion that measured scipy's accuracy rather than ours, and restored the file's left/right/two-sided ordering in batch_util. * Conform to the Array return annotations from #2206 Narrow the new classes' returns from ArrayLike to Array, take Optional keys in sample with the accompanying assert, cast constraint bounds to NumLike, and declare _support at class level — matching what #2206 applied to the existing truncated classes. The generic test_output_is_array added by that PR also caught a broadcasting bug in the bisection helpers: promote_shapes can leave low and high at size-1 dims while the cdf carries the full batch shape, so the fori_loop carry changed shape mid-loop for batched parameters. The helpers now take the target shape explicitly, broadcast from the quantile shape and the distribution's batch shape.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR for issue #969 . Contains initial implementation that performs uniform sampling + inverse CDF of Left/Right/Doubly truncated Gamma. Relies on tensorflow functionality for igammainv function, which is not yet implemented at the lax/jax level (see jax-ml/jax#5350).
There is a test that fails, but it is not clear to me if this is purely a numerical issue with the uniform + iCDF sampling, or a larger issue that I missed at the time I implemented things.