Enable Enzyme gradient tests across the test suite (Julia ≥ 1.12) - #708
Draft
Parvm1102 wants to merge 3 commits into
Draft
Enable Enzyme gradient tests across the test suite (Julia ≥ 1.12) #708Parvm1102 wants to merge 3 commits into
Parvm1102 wants to merge 3 commits into
Conversation
Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
Parvm1102
marked this pull request as draft
August 28, 2026 08:56
Member
|
NNlib should be registered in a few minutes |
`Flux.destructure`'s `re` writes the gradient into the trainable leaves but keeps every other array (e.g. BatchNorm's running stats) at its model value. Zygote returns `nothing` for those fields so `check_equal_leaves` skipped them, but Enzyme returns a full shadow model with the correct zero gradient there and the comparison failed on state that is not a gradient at all. Rebuild the reference from a zeroed model so those leaves compare as zeros. This enables Enzyme on the GNNChain Parallel/BatchNorm test site for non-dense graphs; verified neutral for EdgeConv and GNNlib's propagate. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
Contributor
Author
|
The GNN julia 1 job terminated because the resource limits. Locally the TransformerConv took around 7gb to run, plus other tests might have exhausted the resources |
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.
I am creating a draft PR for enabling enzyme tests after the recent fixes on NNlib.jl, Enzyme.jl and this repo.
It is draft because NNlibv 0.9.45 bump commit is there but hasn't released yet on the julia registry. All CI checks will be failing because of that.
Another reason for being draft PR is because of enzyme, the test suite time has increased drastically. It will now take hours to complete running on CI. Most layers add 3-15 minutes of time each. TransformerConv now takes ~45 minutes, GConvLSTM 33 min, GConvGRU 27 min, GCNConv 22 min and so on. Need to discuss how should we proceed.
Current status for Enzyme:
The Dense failures and DConv will only require simple fixes.