Test against Enzyme / Reactant? #655

Activity

1. Parvm1102 commented on Mar 25, 2026

   

   Parvm1102

   on Mar 25, 2026

   Contributor

   More actions

   @CarloLucibello I can work on this. Should I follow a similar approach to the existing Mooncake testing?
2. wsmoses commented on Jul 2, 2026

   

   wsmoses

   on Jul 2, 2026

   Author

   More actions

   Bumping here
3. wsmoses commented on Jul 18, 2026

   

   wsmoses

   on Jul 18, 2026

   Author

   More actions

   bumping here, @CarloLucibello / @Parvm1102 ?
4. Parvm1102 commented on Jul 26, 2026

   

   Parvm1102

   on Jul 26, 2026

   Contributor

   More actions

   @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 propagate fast paths e_mul_xj / w_mul_xj.

   The following are the failures:

   1. 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.

   2. 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.

   3. EnzymeInternalError on GCNConv's dense/sparse paths (:coo passes).

   4. Compile hang on TGCNCell: The pullback compilation never returns (killed after 50 min).

   5. 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.
5. wsmoses commented on Jul 26, 2026

   

   wsmoses

   on Jul 26, 2026

   Author

   More actions

   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]
6. wsmoses commented on Jul 30, 2026

   

   wsmoses

   on Jul 30, 2026

   Author

   More actions

   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
7. wsmoses commented on Jul 31, 2026

   

   wsmoses

   on Jul 31, 2026

   Author

   More actions

   also bumping reactant too [and should that get its own issue @CarloLucibello ]