Add GNNGraphsMooncakeExt.jl

CHANGELOG.md

@@ -10,7 +10,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and the packages adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 Entries link to the pull request that introduced them.
 
-## GNNGraphs.jl — Unreleased (towards 1.5.2)
+## GNNGraphs.jl — Unreleased (towards 1.6.0)
+
+**Added**
+- 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**
 - 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]).
@@ -230,4 +233,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
 [#695]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/695
 [#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
 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662

GNNGraphs/Project.toml

@@ -1,6 +1,6 @@
 name = "GNNGraphs"
 uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c"
-version = "1.5.2-DEV"
+version = "1.6.0-DEV"
 authors = ["Carlo Lucibello and contributors"]
 
 [workspace]
@@ -23,10 +23,12 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 
 [extensions]
 GNNGraphsCUDAExt = "CUDA"
+GNNGraphsMooncakeExt = "Mooncake"
 GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
 
 [compat]
@@ -38,6 +40,7 @@ KrylovKit = "0.8, 0.9, 0.10"
 LinearAlgebra = "1"
 MLDataDevices = "1.0"
 MLUtils = "0.4"
+Mooncake = "0.5.24"
 NNlib = "0.9"
 NearestNeighbors = "0.4"
 Random = "1"

GNNGraphs/ext/GNNGraphsMooncakeExt.jl

@@ -0,0 +1,16 @@
+module GNNGraphsMooncakeExt
+
+using GNNGraphs
+import Mooncake
+
+# Without these rules Mooncake traces into graph-structure code and fails on CUDA:
+# host→device index copies in `add_self_loops`, `findall` on `CuArray` in
+# `_findnz_idx`, bounds checks in `_edge_values` (Mooncake's CUDA `getindex` rule
+# covers only float/complex arrays). No gradient is lost: all three yield integers,
+# whose Mooncake tangent is `NoTangent` anyway. Float edge values stay on the AD path.
+Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(add_self_loops), GNNGraph}
+Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(GNNGraphs._findnz_idx), Any}
+Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{
+    typeof(GNNGraphs._edge_values), AbstractMatrix{<:Integer}, Any, Any}
+
+end # module

GNNGraphs/src/convert.jl

@@ -80,9 +80,13 @@ end
 
 CRC.@non_differentiable _findnz_idx(A)
 
+# Values of `A` at the extracted edge positions. Linear indexing keeps float edge
+# weights differentiable on GPU, unlike `A[nz]` with CartesianIndex.
+_edge_values(A::AbstractMatrix, s, t) = vec(A)[s .+ (t .- 1) .* size(A, 1)]
+
 function to_coo(A::ADJMAT_T; dir = :out, num_nodes = nothing, weighted = true)
-    s, t, nz = _findnz_idx(A)
-    v = A[nz]
+    s, t, _ = _findnz_idx(A)
+    v = _edge_values(A, s, t)
     if dir == :in
         s, t = t, s
     end
@@ -126,8 +130,8 @@ end
 
 function _to_coo_graph(g::GNNGraph{<:ADJMAT_T})
     A = g.graph
-    s, t, nz = _findnz_idx(A)
-    v = A[nz]
+    s, t, _ = _findnz_idx(A)
+    v = _edge_values(A, s, t)
     return GNNGraph((s, t, v), g.num_nodes, g.num_edges, g.num_graphs,
                     g.graph_indicator, g.ndata, g.edata, g.gdata)
 end

GNNGraphs/test/Project.toml

@@ -8,6 +8,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -26,5 +27,6 @@ GNNGraphs = {path = ".."}
 
 [compat]
 Enzyme = "0.13.197"
+Mooncake = "0.5.24"
 NNlib = "0.9.38"
 StableRNGs = "1.0.4"

GNNGraphs/test/ext/Mooncake.jl

@@ -0,0 +1,79 @@
+@testitem "GNNGraphsMooncakeExt zero-derivative rules" setup=[GraphsTestModule] begin
+    using .GraphsTestModule
+    import Mooncake
+    import NNlib
+
+    # Loading Mooncake triggers the extension.
+    @test Base.get_extension(GNNGraphs, :GNNGraphsMooncakeExt) !== nothing
+
+    # Mooncake is only exercised on Julia >= 1.12, as in the other test suites.
+    if VERSION >= v"1.12"
+        function mooncake_gradient(f, x)
+            cache = Mooncake.prepare_gradient_cache(f, x)
+            _, grads = Mooncake.value_and_gradient!!(cache, f, x)
+            return grads[2]
+        end
+
+        rng = MersenneTwister(17)
+
+        # Gradients through `add_self_loops` (zero-derivative rule).
+        for graph_type in GRAPH_TYPES
+            g = rand_graph(rng, 6, 10; graph_type)
+            x = randn(rng, Float32, 3, g.num_nodes)
+            loss(x) = sum(abs2, NNlib.gather(x, edge_index(add_self_loops(g))[1]))
+            gm = mooncake_gradient(loss, x)
+            @test gm≈gradient(loss, x)[1] rtol=1e-4
+            @test gm≈ngradient(loss, x)[1] rtol=1e-4
+        end
+
+        # Gradients through `edge_index` on adjacency-matrix graphs.
+        for graph_type in (:dense, :sparse)
+            g = rand_graph(rng, 6, 10; graph_type)
+            x = randn(rng, Float32, 3, g.num_nodes)
+            loss(x) = sum(abs2, NNlib.gather(x, edge_index(g)[1]))
+            gm = mooncake_gradient(loss, x)
+            @test gm≈gradient(loss, x)[1] rtol=1e-4
+            @test gm≈ngradient(loss, x)[1] rtol=1e-4
+        end
+
+        # Float adjacency: edge weights stay on the AD path.
+        A = Float32[0 1 0 2; 1 0 3 0; 0 3 0 1; 2 0 1 0]
+        loss_w(A) = sum(abs2, get_edge_weight(GNNGraph(A, graph_type = :dense)))
+        gw = mooncake_gradient(loss_w, copy(A))
+        @test gw≈2 .* A rtol=1e-4
+        @test gw≈gradient(loss_w, A)[1] rtol=1e-4
+    end
+end
+
+@testitem "GNNGraphsMooncakeExt on GPU" setup=[GraphsTestModule] tags=[:gpu] begin
+    using .GraphsTestModule
+    import Mooncake
+    import NNlib
+
+    dev = gpu_device(force = true)
+    # Mooncake's GPU rules are CUDA-only.
+    if VERSION >= v"1.12" && dev isa CUDADevice
+        function mooncake_gradient(f, x)
+            cache = Mooncake.prepare_gradient_cache(f, x)
+            _, grads = Mooncake.value_and_gradient!!(cache, f, x)
+            return grads[2]
+        end
+
+        rng = MersenneTwister(17)
+
+        # Without the rules Mooncake errors here instead of returning a gradient.
+        for graph_type in (:coo, :dense)
+            g = rand_graph(rng, 6, 10; graph_type)
+            x = randn(rng, Float32, 3, g.num_nodes)
+            g_gpu, x_gpu = dev(g), dev(x)
+            loss(x) = sum(abs2, NNlib.gather(x, edge_index(add_self_loops(g))[1]))
+            loss_gpu(x) = sum(abs2, NNlib.gather(x, edge_index(add_self_loops(g_gpu))[1]))
+            @test Array(mooncake_gradient(loss_gpu, x_gpu))≈mooncake_gradient(loss, x) rtol=1e-4
+        end
+
+        # Float adjacency: edge weights stay on the AD path.
+        A = Float32[0 1 0 2; 1 0 3 0; 0 3 0 1; 2 0 1 0]
+        loss_w(A) = sum(abs2, get_edge_weight(GNNGraph(A, graph_type = :dense)))
+        @test Array(mooncake_gradient(loss_w, dev(A)))≈2 .* A rtol=1e-4
+    end
+end