Fix NaN gradients in Hurdle distributions for continuous components#8057
Merged
ricardoV94 merged 3 commits intoJan 17, 2026
Merged
Conversation
The dlogp function was returning NaN for HurdleGamma and HurdleLogNormal when observed data contained zeros. This was caused by pt.where evaluating both branches - even when selecting log(1-psi) for zero values, the logp(dist, 0) computation produced -inf for continuous distributions, resulting in NaN gradients. The fix replaces zero values with a safe dummy value (1.0) before computing the logp of the continuous distribution. Since pt.where still selects the correct branch based on the condition, the final logp result is unchanged, but the gradient computation is now well-defined. Fixes pymc-devs#8053
ricardoV94
approved these changes
Jan 16, 2026
Member
|
This is also fixed by pymc-devs/pytensor#1850 but this implementation is more robust |
Member
|
Pre-commit is failing |
ricardoV94
reviewed
Jan 16, 2026
|
|
||
| check_logp(HurdleLogNormal, Rplus, {"psi": Unit, "mu": R, "sigma": Rplusbig}, logp_fn) | ||
|
|
||
| def test_hurdle_gamma_dlogp_no_nan(self): |
Member
There was a problem hiding this comment.
Can you pytest.mark.parametrize the test on HurdleGamma | HurdleLogNormal ? The logic is nearly the same
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #8057 +/- ##
==========================================
+ Coverage 91.42% 91.44% +0.01%
==========================================
Files 117 117
Lines 19154 19155 +1
==========================================
+ Hits 17512 17516 +4
+ Misses 1642 1639 -3
🚀 New features to boost your workflow:
|
Contributor
Author
|
@ricardoV94 ping |
|
|
Member
|
Thanks @pirzada-ahmadfaraz |
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
2 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.
Summary
Fixes the
dlogpfunction returning NaN forHurdleGammaandHurdleLogNormalwhen observed data contains zeros.Problem
The gradient computation (
dlogp) was failing becausept.whereevaluates both branches before selecting one. When computinglogp(dist, 0)for continuous distributions like Gamma, this produces-inf, which results in NaN gradients - even though that branch is not selected for zero values.Solution
Replace zero values with a safe dummy value (1.0) before computing the logp of the continuous distribution:
Since
pt.wherestill selects the correct branch based on the condition, the final logp result is unchanged, but the gradient computation is now well-defined.Testing
Added regression tests for both
HurdleGammaandHurdleLogNormalthat verifydlogpdoes not return NaN values.Fixes #8053