Skip to content

Evaluate unit-concentration Gamma densities correctly at zero - #361

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/gamma-density-at-zero-20261006
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/gamma-density-at-zero-20261006

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Summary

Fixes #360.

Use jax.scipy.special.xlogy for the concentration-offset times log(value) term. This evaluates the zero-coefficient case correctly at value zero instead of forming 0 * -inf. Unit-concentration Gamma densities now match their exponential special case, including through prob and Independent products. No CDF or sampling code is changed.

Validation

python -m pytest -q distrax/_src/distributions/gamma_test.py distrax/_src/distributions/gamma_boundary_test.py

181 tests and 156 subtests passed; 30 existing variants were skipped. Fourteen new tests include analytic exponential densities, SciPy boundary limits, float32/float64, broadcast concentrations/rates, JIT, an Independent product, rate derivatives and positive-value derivative controls. Ten new cases fail on unchanged upstream and four controls pass.

The existing Gamma equivalence tests also pass; they emitted upstream TensorFlow Probability/JAX deprecation warnings. This patch does not introduce support validation or define observation derivatives at the boundary. GPU/TPU and the full Distrax suite were not run.

New-test formatting, Python syntax checks, scoped static checks and git diff --check pass. No dependency or workflow changes.

Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Gamma log probability returns NaN at the exponential boundary

1 participant