Parvm1102 · Aug 28, 2026 · Aug 28, 2026 · Aug 28, 2026 · Oct 13, 2026 · Oct 13, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -33,7 +33,10 @@ Entries link to the pull request that introduced them.
 **Removed**
 - Removed the `A3TGCN` layer, mirroring its removal from the Flux frontend ([#696]).
 
-## GNNlib.jl — Unreleased (towards 1.4.1)
+## GNNlib.jl — Unreleased (towards 1.5.0)
+
+**Added**
+- Added an `EnzymeCore` package extension with an Enzyme reverse rule for the Chebyshev core of `cheb_conv`, used by `ChebConv`, `GConvGRUCell` and `GConvLSTMCell`. Enzyme previously differentiated the Chebyshev recursion inline at every call, which made compiling the gradient of the recurrent cells several times slower ([#713]).
 
 **Fixed**
 - Fixed `Enzyme.gradient` failing when differentiating `GCNConv`, `SGConv` and `TAGConv` on `:dense`/`:sparse` adjacency graphs: their adjacency-matrix fallbacks now convert via the Enzyme-differentiable `GNNGraphs._to_coo_graph` instead of the keyword `GNNGraph` constructor. Requires GNNGraphs ≥ 1.5.2 and, for Enzyme, Enzyme ≥ 0.13.197 ([#703]).
@@ -238,4 +241,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
 [#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704
 [#707]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/707
 [#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706
+[#713]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/713
 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662
diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml
@@ -1,7 +1,7 @@
 name = "GNNlib"
 uuid = "a6a84749-d869-43f8-aacc-be26a1996e48"
 authors = ["Carlo Lucibello and contributors"]
-version = "1.4.1-DEV"
+version = "1.5.0-DEV"
 
 [workspace]
 projects = ["test", "docs", "../GNNGraphs"]
@@ -19,12 +19,14 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
 GNNlibAMDGPUExt = "AMDGPU"
 GNNlibCUDAExt = "CUDA"
+GNNlibEnzymeCoreExt = "EnzymeCore"
 GNNlibMetalExt = "Metal"
 GNNlibMooncakeExt = "Mooncake"
 
@@ -33,6 +35,7 @@ AMDGPU = "1, 2.7"
 CUDA = "5, 6"
 ChainRulesCore = "1.24"
 DataStructures = "0.18, 0.19"
+EnzymeCore = "0.8.18"
 GNNGraphs = "1.5.2"
 LinearAlgebra = "1"
 Metal = "1.0"

diff --git a/GNNlib/ext/GNNlibEnzymeCoreExt.jl b/GNNlib/ext/GNNlibEnzymeCoreExt.jl
@@ -0,0 +1,82 @@
+module GNNlibEnzymeCoreExt
+
+using GNNlib: GNNlib
+using LinearAlgebra: mul!
+import EnzymeCore
+using EnzymeCore: Const, Duplicated
+using EnzymeCore.EnzymeRules: EnzymeRules, RevConfigWidth, AugmentedReturn, needs_primal,
+                              needs_shadow, overwritten, primal_type, shadow_type,
+                              runtime_activity
+
+# Reverse rule for `GNNlib._cheb_core`, the array part of `cheb_conv`.
+# Without it Enzyme differentiates the Chebyshev recursion inline at every call site, and the
+# recurrent cells call `cheb_conv` up to eight times (`GConvLSTMCell`), which makes their
+# Enzyme compilation several times slower.
+# The scaled Laplacian `L` depends only on the graph and is treated as a constant
+# (`scaled_laplacian` is marked inactive in GNNGraphsEnzymeCoreExt).
+# Batched (width > 1) and other annotations fall back to Enzyme's own differentiation.
+
+const ChebArg = Union{Const, Duplicated}
+
+# Under runtime activity, a `Duplicated` argument whose shadow is its primal is inactive.
+_is_active(config, x) = x isa Duplicated && !(runtime_activity(config) && x.dval === x.val)
+
+# The Chebyshev terms Z₁ = X, Z₂ = X * L, Zₖ = 2 Zₖ₋₁ L - Zₖ₋₂, as in `_cheb_core`.
+function _cheb_terms(W, X, L)
+    Z2 = X * L
+    Zs = [convert(typeof(Z2), X), Z2]
+    for k in 3:size(W, 3)
+        push!(Zs, 2 * Zs[k - 1] * L - Zs[k - 2])
+    end
+    return Zs
+end
+
+function EnzymeRules.augmented_primal(config::RevConfigWidth{1},
+                                      ::Const{typeof(GNNlib._cheb_core)}, ::Type{RT},
+                                      W::ChebArg, b::ChebArg, X::ChebArg,
+                                      L::ChebArg) where {RT}
+    # `overwritten` flags the function itself first, then (W, b, X, L).
+    ow = overwritten(config)
+    Wv = ow[2] ? copy(W.val) : W.val
+    Xv = ow[4] ? copy(X.val) : X.val
+    Lv = ow[5] ? copy(L.val) : L.val
+    Zs = _cheb_terms(Wv, Xv, Lv)
+    Y = view(Wv, :, :, 1) * Zs[1]
+    for k in 2:size(Wv, 3)
+        Y = Y .+ view(Wv, :, :, k) * Zs[k]
+    end
+    Y = Y .+ b.val
+    primal = needs_primal(config) ? Y : nothing
+    shadow = needs_shadow(config) ? zero(Y) : nothing
+    tape = (Zs, shadow, Wv, Lv)
+    return AugmentedReturn{primal_type(config, RT), shadow_type(config, RT), typeof(tape)}(
+        primal, shadow, tape)
+end
+
+function EnzymeRules.reverse(config::RevConfigWidth{1},
+                             ::Const{typeof(GNNlib._cheb_core)}, ::Type{RT}, tape,
+                             W::ChebArg, b::ChebArg, X::ChebArg, L::ChebArg) where {RT}
+    Zs, dY, Wv, Lv = tape
+    dY === nothing && return (nothing, nothing, nothing, nothing)
+    K = length(Zs)
+    if _is_active(config, W)
+        for k in 1:K
+            mul!(view(W.dval, :, :, k), dY, Zs[k]', true, true)
+        end
+    end
+    if _is_active(config, b)
+        b.dval .+= vec(sum(dY; dims = 2))
+    end
+    if _is_active(config, X)
+        dZ = [view(Wv, :, :, k)' * dY for k in 1:K]
+        for k in K:-1:3 # Zₖ = 2 Zₖ₋₁ L - Zₖ₋₂
+            dZ[k - 1] .+= 2 .* (dZ[k] * Lv')
+            dZ[k - 2] .-= dZ[k]
+        end
+        dZ[1] .+= dZ[2] * Lv' # Z₂ = Z₁ L
+        X.dval .+= dZ[1]
+    end
+    return (nothing, nothing, nothing, nothing)
+end
+
+end # module
diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl
@@ -85,16 +85,22 @@ function cheb_conv(l, g::GNNGraph, X::AbstractMatrix{T}) where {T}
     @assert size(X, 1) == size(l.weight, 2) "Input feature size must match input channel size."
 
     L̃ = scaled_laplacian(g, eltype(X))
+    return _cheb_core(l.weight, l.bias, X, L̃)
+end
 
+# Array core of `cheb_conv`: `Y = Σₖ W[:, :, k] * Zₖ .+ b` over the Chebyshev terms
+# `Z₁ = X`, `Z₂ = X * L`, `Zₖ = 2 Zₖ₋₁ L - Zₖ₋₂`. Kept separate so that
+# `GNNlibEnzymeCoreExt` can give Enzyme a reverse rule for it.
+function _cheb_core(W::AbstractArray{<:Any, 3}, b, X::AbstractMatrix, L::AbstractMatrix)
     Z_prev = X
-    Z = X * L̃
-    Y = view(l.weight, :, :, 1) * Z_prev
-    Y = Y .+ view(l.weight, :, :, 2) * Z
-    for k in 3:(l.k)
-        Z, Z_prev = 2 * Z * L̃ - Z_prev, Z
-        Y = Y .+ view(l.weight, :, :, k) * Z
-    end
-    return Y .+ l.bias
+    Z = X * L
+    Y = view(W, :, :, 1) * Z_prev
+    Y = Y .+ view(W, :, :, 2) * Z
+    for k in 3:size(W, 3)
+        Z, Z_prev = 2 * Z * L - Z_prev, Z
+        Y = Y .+ view(W, :, :, k) * Z
+    end
+    return Y .+ b
 end
 
 ####################### GraphConv ######################################

diff --git a/GNNlib/test/Project.toml b/GNNlib/test/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -22,4 +23,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 GNNlib = {path = ".."}
 
 [compat]
-NNlib = "0.9.38"
+NNlib = "0.9.45"
diff --git a/GNNlib/test/msgpass.jl b/GNNlib/test/msgpass.jl
@@ -128,11 +128,19 @@ end
 
 @testitem "propagate" setup=[TestModuleGNNlib] begin
     using .TestModuleGNNlib
+    using Flux: Flux
 
     @testset "copy_xj +" begin
         for g in TEST_GRAPHS
             f(g, x) = propagate(copy_xj, g, +, xj = x)
-            test_gradients(f, g, g.x; test_grad_f=false)
+            # On :dense graphs the `copy_xj`/`+` specialization multiplies by
+            # `adjacency_matrix(g, T; weighted = false)`, whose result is union-typed
+            # (`to_dense` rebinds `A` across its branches), and Enzyme fails there
+            # with an IllegalTypeAnalysisException. The other storage types are fine.
+            test_gradients(f, g, g.x; test_grad_f=false,
+                           ad_backends = get_graph_type(g) == :dense ?
+                                         [Flux.AutoZygote(), Flux.AutoMooncake()] :
+                                         [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()])
         end
     end
 

diff --git a/GNNlib/test/test_module.jl b/GNNlib/test/test_module.jl
@@ -37,7 +37,7 @@ import Reexport: @reexport
 @reexport using SparseArrays
 @reexport using Test, Random, Statistics
 @reexport using MLDataDevices
-using Functors: fmapstructure_with_path
+using Functors: fmap, fmapstructure_with_path
 using FiniteDifferences: FiniteDifferences
 using Zygote: Zygote
 using Flux: Flux
@@ -48,10 +48,18 @@ if TEST_MOONCAKE
     import Mooncake
 end
 
+# Enzyme.jl requires Julia >= 1.12 for the graph code paths exercised here:
+# on 1.10 it recurses forever through `fieldnames(::UnionAll)` in GNNGraph's
+# property sugar (https://github.com/EnzymeAD/Enzyme.jl/issues/3423).
+const TEST_ENZYME = VERSION >= v"1.12"
+if TEST_ENZYME
+    import Enzyme # loads FluxEnzymeExt, which backs Flux.AutoEnzyme()
+end
+
 # from this module
 export D_IN, D_OUT, GRAPH_TYPES, TEST_GRAPHS,
-       test_gradients, finitediff_withgradient, 
-       check_equal_leaves, gpu_backend, TEST_MOONCAKE
+       test_gradients, finitediff_withgradient,
+       check_equal_leaves, gpu_backend, TEST_MOONCAKE, TEST_ENZYME
 
 
 const D_IN = 3
@@ -95,7 +103,8 @@ function test_gradients(
             test_grad_x = true,
             # Reference AD: finite differences for CPU tests, Zygote (on CPU) for GPU tests.
             reference = test_gpu ? Flux.AutoZygote() : :finitediff,
-            ad_backends = test_gpu ? [] : [Flux.AutoZygote(), Flux.AutoMooncake()],
+            ad_backends = test_gpu ? [] :
+                          [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()],
             loss = (f, g, xs...) -> mean(f(g, xs...)),
             )
 
@@ -104,6 +113,11 @@ function test_gradients(
         ad_backends = filter(b -> !(b isa Flux.AutoMooncake), ad_backends)
     end
 
+    # Enzyme requires Julia >= 1.12. Layers it cannot handle opt out per call site.
+    if !TEST_ENZYME
+        ad_backends = filter(b -> !(b isa Flux.AutoEnzyme), ad_backends)
+    end
+
     ## Let's make sure first that the forward pass works.
     l = loss(f, graph, xs...)
     @assert l isa Number
@@ -149,7 +163,10 @@ function test_gradients(
             f64 = f |> Flux.f64
             ps, re = Flux.destructure(f64)
             y_fd, g_fd = finitediff_withgradient(ps -> loss(re(ps), graph, xs...), ps)
-            y_fd, (re(g_fd[1]),)
+            # Rebuild from a zeroed model, so non-trainable arrays (e.g. BatchNorm's
+            # running stats) compare as zero gradients instead of as model values.
+            _, re0 = Flux.destructure(fmap(x -> x isa AbstractArray ? zero(x) : x, f64))
+            y_fd, (re0(g_fd[1]),)
         else
             Flux.withgradient(f -> loss(f, graph, xs...), reference, f)
         end

diff --git a/GraphNeuralNetworks/test/Project.toml b/GraphNeuralNetworks/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -9,6 +10,7 @@ GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -19,3 +21,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [sources]
 GraphNeuralNetworks = {path = ".."}
+
+[compat]
+NNlib = "0.9.45"
diff --git a/GraphNeuralNetworks/test/layers/basic.jl b/GraphNeuralNetworks/test/layers/basic.jl
@@ -18,7 +18,11 @@
 
         Flux.testmode!(gnn)
 
-        test_gradients(gnn, g, x, rtol = 1e-5, ad_backends = [Flux.AutoZygote()])
+        # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where
+        # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix.
+        test_gradients(gnn, g, x, rtol = 1e-5,
+                       ad_backends = GRAPH_T == :dense ? [Flux.AutoZygote()] :
+                                     [Flux.AutoZygote(), Flux.AutoEnzyme()])
 
         @testset "constructor with names" begin
             m = GNNChain(GCNConv(din => d),
@@ -53,7 +57,10 @@
 
             Flux.trainmode!(gnn)
 
-            test_gradients(gnn, g, x, rtol = 1e-4, atol=1e-4, ad_backends = [Flux.AutoZygote()])
+            # Same :dense limitation as above; Enzyme handles the other graph types.
+            test_gradients(gnn, g, x, rtol = 1e-4, atol=1e-4,
+                           ad_backends = GRAPH_T == :dense ? [Flux.AutoZygote()] :
+                                         [Flux.AutoZygote(), Flux.AutoEnzyme()])
         end
     end