You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
I ran the full gradient test suite through Enzyme (Flux.withgradient(f, AutoEnzyme(), ...), Flux 0.16.10 / Enzyme 0.13.195 / NNlib 0.9.38 / Julia 1.12.5), testing every layer against finite differences over all three graph storages (:coo, :dense, :sparse). Summary of the results:
Passing everywhere (now enabled in the test suite): EdgeConv, NNConv, ResGatedGraphConv, CGConv, GlobalPool, DCGRUCell, EvolveGCNOCell, and the propagate fast paths e_mul_xj / w_mul_xj.
The following are the failures:
NNlib's scatter! Enzyme rule lacks reverse for mean/max: augmented_primal accepts any op::Const, but reverse is only defined for +/-, so Enzyme commits to the rule and then fails:
MethodError: no method matching reverse(::RevConfigWidth{...}, ::Const{typeof(NNlib.scatter!)},
::Type{DuplicatedNoNeed{Matrix{Float32}}}, ::Vector{Int64}, ::Const{typeof(Statistics.mean)}, ...)
This causes following layers to fail:
GATConv
GATv2Conv
TransformerConv
AGNNConv
GlobalAttentionPool
GINConv, GraphConv(aggr=mean), copy_xj+mean).
NNlib already has these pullbacks for ChainRules (∇scatter), so extending the Enzyme rule looks mostly like reuse.
IllegalTypeAnalysisException on layers that build graphs/adjacencies inside their forward
ChebConv
DConv
SGConv
TAGConv
the GNNRecurrence temporal layers, dense-storage paths of GCNConv/GraphConv/SAGEConv): traces point at the Union-heavy GNNGraph keyword constructor. A GNNlibEnzymeExt with rules for the propagate fast paths (analogous to GNNlibMooncakeExt) could sidestep this.
EnzymeInternalError on GCNConv's dense/sparse paths (:coo passes).
Compile hang on TGCNCell: The pullback compilation never returns (killed after 50 min).
Wrong gradients on MEGNetConv and EvolveGCNOCell-with-state Enzyme runs but disagrees with finite differences, while Zygote and Mooncake agree with finitediff on the identical calls. Possibly aliasing-related (inputs also reachable through the Const-captured graph). I tried running them multiple times, results were same.
The custom rule ones make sense and need rules [which I presume you'll be able to add]. The 3/4/5 ones, can you open mwe's for [ideally as simple/few packages as possible]
bumping back here, as far as I can tell the answer to most of these is the need for some rules [at let me know if that doesn't suffice], or the check was wrong due to aliasing
@wsmoses we are stuck with #708 since Enzyme compilation of gradients takes too much time. Would it be helpful to identify the strugglers? (assuming they are not uniformly slow...)
Activity
Parvm1102 commentedon Mar 25, 2026
@CarloLucibello I can work on this. Should I follow a similar approach to the existing Mooncake testing?
wsmoses commentedon Jul 2, 2026
Bumping here
wsmoses commentedon Jul 18, 2026
bumping here, @CarloLucibello / @Parvm1102 ?
Parvm1102 commentedon Jul 26, 2026
@wsmoses @CarloLucibello
I ran the full gradient test suite through Enzyme (
Flux.withgradient(f, AutoEnzyme(), ...), Flux 0.16.10 / Enzyme 0.13.195 / NNlib 0.9.38 / Julia 1.12.5), testing every layer against finite differences over all three graph storages (:coo,:dense,:sparse). Summary of the results:Passing everywhere (now enabled in the test suite): EdgeConv, NNConv, ResGatedGraphConv, CGConv, GlobalPool, DCGRUCell, EvolveGCNOCell, and the
propagatefast pathse_mul_xj/w_mul_xj.The following are the failures:
NNlib's
scatter!Enzyme rule lacksreverseformean/max:augmented_primalaccepts anyop::Const, butreverseis only defined for+/-, so Enzyme commits to the rule and then fails:MethodError: no method matching reverse(::RevConfigWidth{...}, ::Const{typeof(NNlib.scatter!)}, ::Type{DuplicatedNoNeed{Matrix{Float32}}}, ::Vector{Int64}, ::Const{typeof(Statistics.mean)}, ...)This causes following layers to fail:
GraphConv(aggr=mean),copy_xj+mean).NNlib already has these pullbacks for ChainRules (
∇scatter), so extending the Enzyme rule looks mostly like reuse.IllegalTypeAnalysisExceptionon layers that build graphs/adjacencies inside their forwardthe GNNRecurrence temporal layers, dense-storage paths of GCNConv/GraphConv/SAGEConv): traces point at the Union-heavy
GNNGraphkeyword constructor. AGNNlibEnzymeExtwith rules for thepropagatefast paths (analogous toGNNlibMooncakeExt) could sidestep this.EnzymeInternalErroron GCNConv's dense/sparse paths (:coopasses).Compile hang on TGCNCell: The pullback compilation never returns (killed after 50 min).
Wrong gradients on MEGNetConv and EvolveGCNOCell-with-state Enzyme runs but disagrees with finite differences, while Zygote and Mooncake agree with finitediff on the identical calls. Possibly aliasing-related (inputs also reachable through the
Const-captured graph). I tried running them multiple times, results were same.wsmoses commentedon Jul 26, 2026
The custom rule ones make sense and need rules [which I presume you'll be able to add]. The 3/4/5 ones, can you open mwe's for [ideally as simple/few packages as possible]
wsmoses commentedon Jul 30, 2026
bumping back here, as far as I can tell the answer to most of these is the need for some rules [at let me know if that doesn't suffice], or the check was wrong due to aliasing
wsmoses commentedon Jul 31, 2026
also bumping reactant too [and should that get its own issue @CarloLucibello ]
wsmoses commentedon Oct 13, 2026
bump @CarloLucibello
CarloLucibello commentedon Oct 13, 2026
@wsmoses we are stuck with #708 since Enzyme compilation of gradients takes too much time. Would it be helpful to identify the strugglers? (assuming they are not uniformly slow...)
wsmoses commentedon Oct 13, 2026
do you have any mwe issues open for that? a lot of work has gone into reducing compile time recently
wsmoses commentedon Oct 13, 2026
also reactant should be super fast regardless
wsmoses commentedon Oct 14, 2026
@CarloLucibello from things on that pr, it sounds like you're not stuck, right?