-
Notifications
You must be signed in to change notification settings - Fork 36
Fix Matern Grads #478
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
Fix Matern Grads #478
Conversation
Codecov ReportBase: 68.87% // Head: 93.72% // Increases project coverage by
Additional details and impacted files@@ Coverage Diff @@
## master #478 +/- ##
===========================================
+ Coverage 68.87% 93.72% +24.85%
===========================================
Files 52 52
Lines 1346 1354 +8
===========================================
+ Hits 927 1269 +342
+ Misses 419 85 -334
Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here. ☔ View full report at Codecov. |
Codecov is taking a little while to update. This is ready for review. |
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.
Looks good to me, I just think it would be good to add a test for it that fails on the master branch.
Well I implemented this change to fix a test on master, so is it not sufficient that CI now passes? |
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.
Ah I didn't realize that 😄 Somehow I thought you ran into issues in some external example. If it fixes CI on master that's sufficient, of course 👍
Summary
We've been seeing test failures for AD with the
MaternKernel
. This fixes these.Proposed changes
Modifies the implementation of the
kappa
function for theMaternKernel
to utiliseif ... else ... end
instead ofifelse.
The reason that this fixes the problem is that, ifiszero(d)
, the branch which calls the bessel function etc (rather that the one that just returns1
) has someNaN
s and someInf
s floating around in it.The
ifelse
branch evaluates this branch-with-NaN
s, and when Zygote is performing the reverse pass, thisNaN
somehow gets propagated backwards.Conversely,
if ... else ... end
entirely avoids computing anyNaN
s, thereby avoiding the problem.This has the potential downside of yielding worse performance under Zygote, the
MaternKernel
performance was already quite bad, so I don't view this as a problem (it might not even make it much worse tbh).What alternatives have you considered?
The solution that's defnitely going to work + be performant here would be a custom rule, but I don't have the time, and I don't know that this can be improved without sorting out the
besselk
rrule
anyway.Breaking changes
Not breaking.