-
-
Notifications
You must be signed in to change notification settings - Fork 984
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Use approximate log_beta() in .fit(), .predict() #2502
Conversation
n = total_count.clamp(min=0) | ||
k = value.masked_fill(invalid, 0) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is the crux of the NAN gradient fix
@@ -23,7 +23,7 @@ class ExtendedBinomial(Binomial): | |||
|
|||
def log_prob(self, value): | |||
result = super().log_prob(value) | |||
invalid = ~super().support.check(value) | |||
invalid = (value < 0) | (value > self.total_count) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
curous: why this change?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The new version is a little cheaper. The old integer_interval
version additionally checks value % 1 == 0
, but that is already checked in validation by the above line, and it doesn't affect numerical stability.
Thanks for reviewing! |
Addresses #2426
Blocking #2498 (from which this simpler PR was extracted)
This PR:
ExtendedBetaBinomial
related to nan propagates through backward pass even when not accessed pytorch/pytorch#15506. This may have caused spurious rejections during heuristic initialization.log_beta()
inCompartmentalModel.fit()
and.predict()
now that gradients are safe. This reduces likelihood compute cost by about 14% inSuperspredingSEIRModel
, and will reduce Add overdispersed models to contrib.epidemiology #2498 more.safe_log()
topyro.ops.special
which seems like a better home. (this change is non-breaking sincesafe_log()
has not yet been released).Tested
ExtendedBinomial
andExtendedBetaBinomial