Keep derivatives and imaginary parts at ties in logsumexp - #129
Merged
Merged
Conversation
At exact ties, the reduction operators replaced the difference of the arguments by a constant zero to avoid `Inf - Inf = NaN`. This dropped derivatives in Enzyme and Mooncake and discarded differing imaginary parts of complex inputs. Now the difference is replaced only if it is `NaN`. Adds Enzyme and Mooncake tests in a separate environment `test/ad` with its own CI job. Fixes #128. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Evaluate at a less trivial tie point, check primal values for Enzyme and Mooncake, and test Mooncake in forward mode as well. The tie parameters are passed as an explicit argument since Mooncake's forward mode does not support closures that capture differentiable data. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Use single-argument functions with literal constants instead of a parametrised tie function. They capture no data, so Mooncake's forward mode works with them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
devmotion
marked this pull request as ready for review
September 25, 2026 11:23
tpapp
approved these changes
Sep 25, 2026
tpapp
left a comment
tpapp
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. Thanks for keeping the integration tests separate, much cleaner.
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.
At exact ties,
_logsumexp_onepass_opreplaced the difference of the arguments byzero(...)to avoidInf - Inf = NaN. This dropped derivatives in Enzyme and Mooncake and discarded differing imaginary parts of complex inputs (e.g.logsumexp(z for z in (0.0, im))returned0.69 + 1.0iminstead of0.56 + 0.5im). The difference is now replaced only if it isNaN.Enzyme and Mooncake tests live in a separate environment
test/adwith its own CI job on Julia1, since Mooncake pins supported Julia versions (1.10.0, ourmin-patch, is not supported).Fixes #128.
🤖 Generated with Claude Code