Add TruncatedGamma distributions

[deeb01]

Closes #969. Supersedes #1187, inactive since 2022 — same approach, credit to @quattro.

Summary

• add LeftTruncatedGamma, RightTruncatedGamma and TwoSidedTruncatedGamma, dispatched
  by TruncatedGamma(concentration, rate, *, low=None, high=None), following the structure
  @fehiepsi outlined in Truncated Gamma #969
• add gammainccinv to distributions/util.py, wrapping tfp.math.igammacinv like the
  existing gammaincinv; needed for icdf on the upper-tail branch
• TwoSidedTruncatedGamma normalizes with Q(a, λ·low) − Q(a, λ·high) when the interval
  lies above the median, and switches cdf, icdf, mean and variance to match. The
  lower-tail form F(high) − F(low) is exactly zero in single precision for Gamma(2, 1)
  on [30, 40], which would make log_prob +inf
• moments in closed form: E[X^k] = (a)_k / λ^k · [P(a+k, λ·high) − P(a+k, λ·low)] / Z
• three classes rather than one with infinite bounds, because
  jax.grad(lambda a: gammainc(a, jnp.inf)) is nan and would break the gradient w.r.t.
  concentration for left truncation

Two limits, documented in the docstrings: log_prob returns -inf rather than +inf
once the normalizer underflows, and icdf returns nan past the point where
tfp.math.igammainv saturates, which Gamma.icdf also hits on master. A bisection
fallback would lift the second, here or as a follow-up.

Tests

• 42 tests in test/test_distributions.py covering log_prob against scipy,
  normalization and moments against quadrature, cdf/icdf round trip, KS tests on
  samples, batch shapes, both limits, the dispatcher, gradients and NUTS recovery
• registered in the CONTINUOUS parametrization and in batch_util.vmap_over
• make lint passes, as do test_distributions.py, test_gof.py,
  test_distributions_util.py, test_constraints.py, test_transforms.py and
  test_pickle.py. (test_mcmc.py::test_beta_bernoulli_x64[SA] fails for me on master
  too, unrelated to this change.)

Dependencies

None new; tensorflow_probability is already optional and used by gammaincinv.

[github-actions]

Benchmark report

this PR truncated-gamma at 3f92d929 vs baseline master at eda5d6dc

+ run time:     1 faster
  compile time: unchanged across 32 benchmarks

Significant changes (1)

                         ─────── run time ───────     ───── compile time ─────
  benchmark              baseline  this PR      Δ     baseline  this PR      Δ
──────────────────────────────────────────────────────────────────────────────
+ categorical_log_prob     2.3 ms   2.2 ms  -7.0%      71.2 ms  68.6 ms  -3.8%

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ──────── run time ────────     ────── compile time ─────
  benchmark                      baseline  this PR        Δ     baseline   this PR      Δ
─────────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            3.8 ms   3.7 ms    -0.8%     364.5 ms  369.8 ms  +1.4%
+ categorical_log_prob             2.3 ms   2.2 ms    -7.0%      71.2 ms   68.6 ms  -3.8%
  dirichlet_log_prob               631 µs   582 µs  (-7.7%)     410.7 ms  383.6 ms  -6.6%
  dirichlet_sample                55.9 ms  55.7 ms    -0.5%     785.3 ms  801.1 ms  +2.0%
  gamma_log_prob                   2.1 ms   2.1 ms    +1.0%       2.13 s    1.98 s  -7.2%
  gamma_sample                    25.8 ms  26.2 ms    +1.3%     774.1 ms  762.1 ms  -1.5%
  lkj_cholesky_sample              5.7 ms   5.8 ms    +0.2%       1.13 s    1.13 s  +0.3%
  mixture_same_family_log_prob     974 µs   976 µs    +0.2%      78.5 ms   77.7 ms  -1.0%
  multivariate_normal_log_prob     248 µs   252 µs    +1.6%     574.1 ms  568.0 ms  -1.1%
  normal_log_prob                  533 µs   538 µs    +1.0%      56.4 ms   53.5 ms  -5.2%
  normal_sample                   21.0 ms  20.6 ms    -1.6%     190.0 ms  199.7 ms  +5.1%
  stick_breaking_transform         6.3 ms   6.3 ms    +0.5%     195.8 ms  198.0 ms  +1.1%
  student_t_log_prob               3.1 ms   3.0 ms    -1.6%      76.4 ms   76.0 ms  -0.4%
  truncated_normal_log_prob        586 µs   544 µs  (-7.1%)      52.2 ms   50.5 ms  -3.2%

handlers

                                  ───────── run time ────────     ────── compile time ──────
  benchmark                       baseline   this PR        Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    38.4 ms   37.3 ms    -3.0%       3.81 s    3.70 s   -2.7%
  log_density_hierarchical          3.0 ms    3.0 ms    -1.5%       1.32 s    1.32 s   -0.7%
  nested_handler_stack              1.2 ms    1.2 ms    +1.0%       634 µs    708 µs  +11.7%
  potential_energy_and_grad          25 µs     25 µs    +0.8%     101.1 ms   98.3 ms   -2.7%
  predictive_forward_sampling     695.1 ms  688.2 ms    -1.0%     169.5 ms  157.5 ms   -7.1%
  trace_seeded_model                707 µs    766 µs  (+8.5%)     540.4 ms  542.8 ms   +0.4%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    710.5 ms  717.6 ms  +1.0%       3.27 s   3.27 s  +0.0%
  nuts_dense_mass_funnel       1.11 s    1.13 s  +1.3%       2.60 s   2.59 s  -0.5%
  nuts_eight_schools           1.10 s    1.13 s  +2.3%       2.51 s   2.69 s  +7.5%
  nuts_hierarchical_glm        4.81 s    4.83 s  +0.4%       4.87 s   4.69 s  -3.7%
  nuts_logistic_regression     1.08 s    1.06 s  -1.5%       3.34 s   3.26 s  -2.5%
  nuts_vectorized_chains       2.48 s    2.45 s  -1.4%       2.92 s   2.75 s  -5.8%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 303.4 ms  302.4 ms  -0.3%       3.23 s   3.14 s  -2.8%
  svi_autodiagonalnormal_hierarchical        985.9 ms    1.00 s  +1.6%       5.25 s   5.11 s  -2.5%
  svi_automultivariatenormal_eight_schools   729.4 ms  724.8 ms  -0.6%       4.51 s   4.49 s  -0.4%
  svi_autonormal_logistic                    742.1 ms  743.0 ms  +0.1%       3.54 s   3.51 s  -0.9%
  svi_multi_particle_elbo                      1.41 s    1.45 s  +3.0%       3.58 s   3.41 s  -4.6%
  svi_trace_mean_field_elbo                    1.29 s    1.28 s  -0.9%       5.49 s   5.38 s  -1.9%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	truncated-gamma
commit	eda5d6dc	3f92d929
numpyro	0.21.0	0.21.0
jax	0.11.1	0.11.1
backend	cpu	cpu
python	3.14.7	3.14.7

Runner: Linux-6.17.0-1022-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[juanitorduz]

Overall LGTM :)

I did an assisted review with Claude and here are the findings :)

1. _partial_moment reuses a tail switch chosen at the wrong order

_use_upper_tail is decided from P(α, λ·low) > 0.5, which is exactly right for the normalizer. But _partial_moment(k) evaluates the same difference at α + k, and the comment "the tail switch carries over unchanged" isn't true: Gamma(α+k) has a larger median, so an interval above the median of Gamma(α) can sit deep in the lower tail of Gamma(α+2), where the Q form is the cancelling one.

It's reachable and it produces a negative variance:

d = dist.TruncatedGamma(0.05, 1.0, low=1e-6, high=1e-5)
d.mean      # 3.985e-06   (quad: 3.9885e-06)
d.variance  # -1.588e-11  (quad: +6.320e-12)

The order-2 difference is the culprit — in f32, Q(2.05, λ·low) - Q(2.05, λ·high) evaluates to exactly 0.0 while P(2.05, λ·high) - P(2.05, λ·low) gives 2.6597e-11, correct to six digits. Deciding the branch per shifted order fixes it outright:

def _partial_moment(self, order):
    shifted = self.base_dist.concentration + order
    rate = self.base_dist.rate
    use_upper = gammainc(shifted, rate * self.low) > 0.5
    return jnp.where(use_upper, ...)

I'd also add a moment case with concentration < 1 and a tail interval to _TRUNCATED_GAMMA_CASES — the current five all have α ≥ 0.5 with intervals near the bulk, so none of them exercise this.

2. Ordinary left truncation silently produces NaN samples

The docstring frames the icdf saturation as "the same limit as Gamma.icdf", but Gamma.sample doesn't go through icdf — it uses standard_gamma. So this is a new failure mode in ordinary use, not an inherited one:

dist.TruncatedGamma(2., 1., low=10.).sample(key, (200_000,))  # 4e-05 NaN
dist.TruncatedGamma(2., 1., low=15.).sample(key, (200_000,))  # 0.6%  NaN

Gamma(2, 1) truncated at 15 is not an exotic object, and 0.6% NaNs flowing into a posterior predictive is the kind of thing that costs someone an afternoon. test_truncated_gamma_saturated_icdf_is_nan currently codifies this as intended behaviour, which I don't think we want to ship. I'd rather see the bisection fallback land in this PR than as a follow-up — a lax.fori_loop bisection on log Q over the tail is maybe 25 lines and vectorises fine. Failing that, the .. note:: needs to say plainly that sample returns NaN in the far tail, and the sampling limit belongs in the class docstring body, not only the note.

Worth fixing, lower priority

variance cancellation in f32. Independent of point 1, E[X²] - E[X]² loses everything when the interval is narrow relative to its location:

d = dist.TruncatedGamma(2.0, 1e-4, low=100.0, high=101.0)
d.variance  # f32: 0.1406   x64: 0.08333  (near-uniform, 1/12)

70% off in the default dtype. I don't think there's a clean closed form for the central moment here, so a docstring note is probably the honest answer — but test_mean_var's tolerance is loose enough to hide this, so it should at least be written down.

mean and variance aren't guarded for the underflowed normalizer. log_prob returns -inf and cdf returns nan, both deliberate, but mean divides by zero and returns nan with no comment. Worth making the three consistent, or documenting why they differ. Relatedly: log_prob → -inf versus cdf → nan for the same underflow is a slightly odd pair; a one-line comment explaining that a density that's identically zero has a defined log-density but no defined conditional CDF would help future readers.

Support closed at zero. RightTruncatedGamma and TwoSidedTruncatedGamma use constraints.interval(0., high), so support(0.) is True and log_prob(0.) is +inf for concentration < 1. The docstring appeals to Beta, but Gamma's own support is positive, i.e. open. constraints.open_interval subclasses _Interval so biject_to still works. Matching the base distribution seems safer than matching Beta here.

API consistency. dist.TruncatedDistribution(dist.Gamma(2., 1.), low=1.) raises a bare AssertionError with no message. Since TruncatedGamma now exists, either dispatch Gamma from TruncatedDistribution to the new classes, or give that assert a message pointing at TruncatedGamma.

Nits

• Broken cross-reference in RightTruncatedGamma's docstring: :class:`~numpyro.distributions.\ncontinuous.Beta` — reST normalises the newline to a space, so the target becomes numpyro.distributions. continuous.Beta and the link won't resolve. The same construction appears in TwoSidedTruncatedGamma.
• test_truncated_gamma_far_tail_normalizer's final assertion (abs(naive - exact) / exact > 1e-10) is testing scipy's f64 accuracy, not our code. For (20, 25) the margin is 1.57e-9 — only ~15x — so it's fragile against a scipy implementation change. I'd drop it; the assert_allclose above it already carries the test's real content.
• reparametrized_params is ["low", "high"], but sampling is pure inverse-CDF and I confirmed jax.grad w.r.t. concentration through both icdf and sample is finite and correct (0.344 and 0.580 for the doctest case). Adding concentration/rate would enable pathwise gradients in SVI. It'd be a departure from how the existing truncated classes are set up, though, so I'd take it or leave it depending on what @fehiepsi prefers.
• gammainccinv has no direct unit test — consistent with gammaincinv, which doesn't either, so purely optional. A round-trip against gammaincc in test_distributions_util.py would be cheap.
• In batch_util.py the new registrations go two-sided → left → right; the file's existing order is left → right → two-sided. Cosmetic.

One question

Have you checked how these behave under TransformReparam / numpyro.infer.reparam, and whether Gamma's KL registrations need a companion? Not blocking, but it'd be good to know the gap is intentional.

[juanitorduz]

Thanks! @deeb01