Add GNNGraphsMooncakeExt.jl

[Parvm1102]

Part of #702 (Mooncake on CUDA tracker).

Mooncake does not use the ChainRules @non_differentiable annotations that keep Zygote working, so it traces into graph-structure functions and fails on CUDA (CPU to GPU index copies, findall on CuArray, bounds checks of integer indexing). This PR gives Mooncake the equivalent rules.

Changes

• New extension GNNGraphs/ext/GNNGraphsMooncakeExt.jl with three Mooncake.@zero_derivative rules, all on integer structure extraction:
  • add_self_loops(::GNNGraph), mirroring the existing @non_differentiable in transform.jl
  • _findnz_idx(::Any), mirroring the existing @non_differentiable in convert.jl
  • _edge_values(::AbstractMatrix{<:Integer}, ...), integer edge values have no derivative
• Rewrote the edge weight extraction in to_coo and _to_coo_graph. A new helper _edge_values(A, s, t) = vec(A)[s .+ (t .- 1) .* size(A, 1)] replaces v = A[nz]. Mooncake has a CUDA rule for linear indexing but not for CartesianIndex indexing. With this, both conversions are ordinary traceable code and float edge weights stay differentiable under Mooncake, same as Zygote and Enzyme.
• Tests in GNNGraphs/test/ext/Mooncake.jl: Mooncake vs Zygote vs finite differences through add_self_loops, edge_index on dense and sparse graphs, _to_coo_graph, and a float-weighted dense adjacency. Mooncake parts only run on Julia >= 1.12, as in the other suites.
• Mooncake weakdep with compat "0.5.24", version bump to 1.6.0-DEV.

This is simpler than the sketch in #702. The original 4 rules marked all of to_coo zero-derivative, which silently dropped edge weight gradients for float dense adjacencies. With the rewrite those rules are not needed and that problem is gone. I checked each remaining rule is really required by probing the bare functions on GPU with the extension disabled.

Tests

CPU (Julia 1.12.5, Mooncake 0.5.45, Enzyme 0.13.198): new testitem 12/12, full GNNGraphs suite 1346 pass and 0 fail (includes the #703 Enzyme test on the rewritten helper), GraphNeuralNetworks GCNConv/GraphConv/SGConv/TAGConv/EdgeConv items 316/316, GNNlib propagate 40/40.

GPU (RTX 3050, CUDA.jl 6.2.1), all of these errored before and now pass, gradients compared against Zygote:

site	max diff vs Zygote
GCNConv coo with self loops, wrt x and params	0.0
GCNConv dense	0.0
SGConv coo, wrt x and params	0.0
EdgeConv, NNConv, CGConv, MEGNetConv, ResGatedGraphConv on dense	<= 1.5e-8
EdgeConv on float-weighted dense	0.0
GNNlib e_mul_xj (matrix e) on dense	0.0
GNNlib w_mul_xj on dense	matches the CPU reference exactly, on a case where Zygote GPU is wrong (the broken = true test in msgpass.jl)

Still failing on GPU, all upstream Mooncake gaps tracked in #702: GAT, GATv2, Transformer, GMMConv, unweighted dense SGConv (missing sum(x; dims) rule), TAGConv and Transformer (CuMatrix + CuMatrix has no rule), TGCN family (repeat has no rule).

[Parvm1102]

@CarloLucibello , please review this pr

[CarloLucibello]

why the integer case was made with zero derivatives?
Does also zygote return zero derivatives for this case?

[Parvm1102]

No, Zygote returns non zero gradients here. Array{<:Integer} is NoTangent anyway, so the rule loses nothing. It is needed because Mooncake's CUDA getindex rule covers just float arrays, so the integer path fails on GPU.

[CarloLucibello]

Array{<:Integer} is NoTangent anyway,

what do you mean by this? Also, the zygote gradient for integer input is not necessarily an integer.

I don't like zygote and mooncake derivatives giving different results, but if there are no better ways to do it then let's just do it and move on.