diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e412e143..d48bfb526 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml index 015e9a68c..a66d48ecc 100644 --- a/GNNGraphs/Project.toml +++ b/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" diff --git a/GNNGraphs/ext/GNNGraphsEnzymeCoreExt.jl b/GNNGraphs/ext/GNNGraphsEnzymeCoreExt.jl new file mode 100644 index 000000000..0dd93cdb9 --- /dev/null +++ b/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 diff --git a/GNNGraphs/test/enzyme.jl b/GNNGraphs/test/enzyme.jl index fd0efd8f6..9b4a5e08e 100644 --- a/GNNGraphs/test/enzyme.jl +++ b/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