Add GNNGraphsEnzymeCoreExt marking scaled_laplacian inactive for Enzyme

CHANGELOG.md

@@ -16,6 +16,7 @@ Entries link to the pull request that introduced them.
 - Added a `Mooncake` package extension, so that Mooncake can differentiate on CUDA through `add_self_loops` and through adjacency-matrix graphs. Float edge weights of adjacency-matrix graphs stay differentiable ([#704]).
 
 **Fixed**
+- Added an `EnzymeCore` package extension marking `scaled_laplacian` inactive for Enzyme, matching its existing `@non_differentiable` ChainRules declaration. Enzyme previously differentiated the Krylov eigensolve inside it and failed; this unblocks `ChebConv`, `GConvGRUCell` and `GConvLSTMCell` ([#706]).
 - Enzyme can now differentiate through the adjacency-matrix → COO graph conversion: a new internal keyword-free helper `_to_coo_graph` avoids the union-typed keyword handling of the `GNNGraph(g; graph_type)` constructor and of `to_coo`, which Enzyme's type analysis cannot compile. Together with upstream fixes in Enzyme ≥ 0.13.197 this removes the `EnzymeInternalError` crash on `:dense`/`:sparse` graphs ([#703]).
 
 ## GNNLux.jl — Unreleased (towards 0.2.0)
@@ -234,4 +235,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
 [#696]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/696
 [#703]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/703
 [#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704
+[#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706
 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662

GNNGraphs/Project.toml

@@ -23,17 +23,20 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 
 [extensions]
 GNNGraphsCUDAExt = "CUDA"
+GNNGraphsEnzymeCoreExt = "EnzymeCore"
 GNNGraphsMooncakeExt = "Mooncake"
 GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
 
 [compat]
 CUDA = "5, 6"
 ChainRulesCore = "1"
+EnzymeCore = "0.8"
 Functors = "0.5"
 Graphs = "1.4"
 KrylovKit = "0.8, 0.9, 0.10"

GNNGraphs/ext/GNNGraphsEnzymeCoreExt.jl

@@ -0,0 +1,11 @@
+module GNNGraphsEnzymeCoreExt
+
+using GNNGraphs
+import EnzymeCore
+
+# `scaled_laplacian` is `@non_differentiable` for ChainRules (see query.jl), but Enzyme
+# does not read those declarations and would differentiate its Krylov eigensolve.
+# It depends only on the graph, never on the node features, so no gradient is lost.
+EnzymeCore.EnzymeRules.inactive(::typeof(GNNGraphs.scaled_laplacian), args...) = nothing
+
+end # module

GNNGraphs/test/enzyme.jl

@@ -23,3 +23,21 @@
     Enzyme.autodiff(Reverse, loss_w, Active, Duplicated(g, dg))
     @test nonzeros(dg.graph) ≈ 2 .* nonzeros(S)
 end
+
+@testitem "Enzyme treats scaled_laplacian as constant" setup=[GraphsTestModule] begin
+    using .GraphsTestModule
+    using Enzyme: Enzyme, Const, Reverse
+
+    # Without GNNGraphsEnzymeCoreExt Enzyme differentiates the Krylov eigensolve
+    # inside `scaled_laplacian` and fails; ChainRules already treats it as constant.
+    # Bidirected cycle: symmetric and without isolated nodes.
+    src = [1, 2, 3, 4, 5, 6, 7, 8]
+    dst = [2, 3, 4, 5, 6, 7, 8, 1]
+    g = GNNGraph(vcat(src, dst), vcat(dst, src))
+    x = randn(MersenneTwister(17), Float32, 3, g.num_nodes)
+    loss(x, g) = sum(abs2, x * scaled_laplacian(g, Float32))
+
+    grad = Enzyme.gradient(Reverse, loss, x, Const(g))[1]
+    gfd = ngradient(x -> loss(x, g), x)[1]
+    @test grad ≈ gfd rtol = 1e-4
+end