Fix Zygote gradient through add_edges

.gitignore

@@ -6,6 +6,8 @@
 *.swo
 Manifest.toml
 .vscode
+.tools/
+.julia-depot/
 LocalPreferences.toml
 .DS_Store
 test.jl

CHANGELOG.md

@@ -10,6 +10,11 @@ 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 1.6.1-DEV
+
+**Fixed**
+- Fixed Zygote differentiation through `add_edges` with non-empty edge metadata by iterating over a collected vector of metadata keys. `sample_neighbors(...; replace=true)` now skips zero-degree nodes instead of trying to sample from an empty neighbor list ([#729]).
+
 ## GNNGraphs.jl 1.6.0 - 2026-11-04
 
 **Added**
@@ -250,5 +255,6 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
 [#717]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/717
 [#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722
 [#718]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/718
+[#729]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/729
 [#728]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/728
 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662

DEVELOPMENT.md

@@ -0,0 +1,27 @@
+# Local development setup
+
+Run PowerShell from the repository root:
+
+```powershell
+.\scripts\setup-dev.ps1
+```
+
+The setup installs Julia 1.13.3 as a portable copy under `.tools`, verifies its
+SHA-256 checksum, then installs the workspace, package, test, docs, and
+JuliaFormatter environments. Julia packages, compiled caches, and artifacts
+stay in `.julia-depot`. Neither setup nor the launcher edits the machine PATH or
+uses the global Julia depot. Both directories are ignored by Git.
+
+Use the repo-local Julia launcher for subsequent work:
+
+```powershell
+.\scripts\julia.ps1 --version
+.\scripts\julia.ps1 --project=GNNGraphs -e 'using Pkg; Pkg.test("GNNGraphs")'
+.\scripts\julia.ps1 --project=GraphNeuralNetworks -e 'using Pkg; Pkg.test("GraphNeuralNetworks")'
+.\scripts\julia.ps1 --project=GNNGraphs/docs GNNGraphs/docs/make.jl
+.\scripts\julia.ps1 --project=.tools/devtools -e 'using JuliaFormatter; format(".")'
+```
+
+The full package test environments are installed, but setup does not execute
+the test suites. GPU backends remain opt-in; install and select the backend
+required by the machine before running GPU-tagged tests.

GNNGraphs/Project.toml

@@ -1,6 +1,6 @@
 name = "GNNGraphs"
 uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c"
-version = "1.6.0"
+version = "1.6.1-DEV"
 authors = ["Carlo Lucibello and contributors"]
 
 [workspace]

GNNGraphs/src/datastore.jl

@@ -172,7 +172,11 @@ function cat_features(ds1::DataStore, ds2::DataStore)
     n1, n2 = getn(ds1), getn(ds2)
     n1 = n1 >= 0 ? n1 : 1
     n2 = n2 >= 0 ? n2 : 1
-    return DataStore(n1 + n2, cat_features(getdata(ds1), getdata(ds2)))
+    data1, data2 = getdata(ds1), getdata(ds2)
+    if isempty(data1) && isempty(data2)
+        return CRC.ignore_derivatives(() -> DataStore(n1 + n2, Dict{Symbol, Any}()))
+    end
+    return DataStore(n1 + n2, cat_features(data1, data2))
 end
 
 function cat_features(dss::AbstractVector{DataStore}; kws...)

GNNGraphs/src/sampling.jl

@@ -177,6 +177,7 @@ function sample_neighbors(g::GNNGraph{<:COO_T}, nodes, K = -1;
     @assert dir ∈ (:in, :out)
     _, eidlist = adjacency_list(g, nodes; dir, with_eid = true)
     for i in 1:length(eidlist)
+        isempty(eidlist[i]) && continue
         if replace
             k = K > 0 ? K : length(eidlist[i])
         else

GNNGraphs/src/utils.jl

@@ -65,10 +65,13 @@ function cat_features(x1::NamedTuple, x2::NamedTuple)
 end
 
 function cat_features(x1::Dict{Symbol, T}, x2::Dict{Symbol, T}) where {T}
-    sort(collect(keys(x1))) == sort(collect(keys(x2))) ||
+    keys1, keys2 = CRC.ignore_derivatives() do
+        sort(collect(keys(x1))), sort(collect(keys(x2)))
+    end
+    keys1 == keys2 ||
         @error "cannot concatenate feature data with different keys"
 
-    return Dict{Symbol, T}([k => cat_features(x1[k], x2[k]) for k in keys(x1)]...)
+    return Dict{Symbol, T}([k => cat_features(x1[k], x2[k]) for k in keys1]...)
 end
 
 function cat_features(x::Dict)
@@ -391,4 +394,4 @@ function color_refinement(g::GNNGraph, x0::AbstractVector{<:Integer})
     end
     num_colors = length(union(x))
     return x, num_colors, niters
-end
+end

GNNGraphs/test/sampling.jl

@@ -25,9 +25,9 @@
         g = rand_graph(10, 40, bidirected = false, graph_type = GRAPH_T)
         sg = sample_neighbors(g, nodes, K; dir, replace = true)
         @test sg.num_nodes == 10
-        @test sg.num_edges == sum(K for i in nodes)
-        @test size(sg.edata.EID) == (sg.num_edges,)
         adjlist = adjacency_list(g; dir)
+        @test sg.num_edges == sum(isempty(adjlist[i]) ? 0 : K for i in nodes)
+        @test size(sg.edata.EID) == (sg.num_edges,)
         s, t = edge_index(sg)
         @test all(s .∈ Ref(nodes))
         for i in nodes
@@ -47,6 +47,24 @@
         @test sg.ndata.x1 == g.ndata.x1[sg.ndata.NID]
         @test length(union(sg.ndata.NID)) == length(sg.ndata.NID)
     end
+
+    @testset "zero-degree nodes with replacement" begin
+        g = GNNGraph(([1], [2]); num_nodes = 3)
+
+        for (dir, node) in ((:in, 1), (:out, 2))
+            sg = sample_neighbors(g, [node], 2; dir, replace = true)
+            @test sg.num_nodes == 3
+            @test sg.num_edges == 0
+            @test isempty(sg.edata.EID)
+            @test edge_index(sg) == (Int[], Int[])
+        end
+
+        for (dir, node) in ((:in, 2), (:out, 1))
+            sg = sample_neighbors(g, [node], 2; dir, replace = true)
+            @test sg.num_edges == 2
+            @test sg.edata.EID == [1, 1]
+        end
+    end
 end
 
 @testitem "induced_subgraph" setup=[GraphsTestModule] begin

GNNGraphs/test/transform.jl

@@ -295,6 +295,38 @@ end
                 gnew = add_edges(g, (snew, tnew, wnew))
                 @test get_edge_weight(gnew) == [w; wnew]
             end
+
+            @testset "differentiate through added edge weights" begin
+                s = [1, 2]
+                t = [2, 3]
+                w = Float32[0.5, 0.7]
+                g = GNNGraph((s, t, w))
+                snew = [1, 2]
+                tnew = [2, 3]
+                wnew = Float32[0.1, 0.2]
+
+                loss_fn(wnew) = sum(get_edge_weight(add_edges(g, (snew, tnew, wnew))))
+                gnew = add_edges(g, (snew, tnew, wnew))
+
+                @test gnew.num_edges == 4
+                @test edge_index(gnew) == ([1, 2, 1, 2], [2, 3, 2, 3])
+                @test get_edge_weight(gnew) == Float32[0.5, 0.7, 0.1, 0.2]
+                @test gradient(loss_fn, wnew)[1] == ones(Float32, 2)
+            end
+
+            @testset "differentiate through added edge data" begin
+                s = [1, 2]
+                t = [2, 3]
+                edata = (x = reshape(Float32[0.5, 0.7], 1, :),)
+                g = GNNGraph((s, t); edata)
+                snew = [1, 2]
+                tnew = [2, 3]
+                xnew = reshape(Float32[0.1, 0.2], 1, :)
+
+                loss_fn(xnew) = sum(add_edges(g, (snew, tnew); edata = (x = xnew,)).edata.x)
+
+                @test gradient(loss_fn, xnew)[1] == ones(Float32, size(xnew))
+            end
         end
     end
 end

GNNLux/reproduce_issue_732.jl

@@ -0,0 +1,27 @@
+using Random
+using GNNLux
+using Lux
+using LuxCore
+
+# Small directed graph with two nodes and two edges.
+g = GNNGraph([1, 2], [2, 1])
+in_dims, out_dims = 3, 4
+x = randn(MersenneTwister(1), Float32, in_dims, g.num_nodes)
+e = randn(MersenneTwister(2), Float32, in_dims, g.num_edges)
+
+# MEGNet's edge update receives source node, destination node, and edge
+# features (3 * in_dims); BatchNorm keeps running statistics in its state.
+ϕe = Chain(Dense(3in_dims => out_dims), BatchNorm(out_dims))
+ϕv = Chain(Dense(in_dims + out_dims => out_dims), Dense(out_dims => out_dims))
+l = MEGNetConv(in_dims, out_dims, ϕe, ϕv)
+
+ps, st = LuxCore.setup(MersenneTwister(3), l)
+edge_state_before = st.ϕe.layer_2
+(x_new, e_new), st_new = l(g, x, e, ps, st)
+edge_state_after = st_new.ϕe.layer_2
+
+println("Output shapes: nodes=$(size(x_new)), edges=$(size(e_new))")
+println("Input BatchNorm running_mean:  ", edge_state_before.running_mean)
+println("Returned BatchNorm running_mean: ", edge_state_after.running_mean)
+println("State changed: ", edge_state_after.running_mean != edge_state_before.running_mean)
+println("Returned full state equals input state: ", st_new == st)

GNNLux/src/layers/conv.jl

@@ -1541,9 +1541,11 @@ end
 
 function (l::MEGNetConv)(g, x, e, ps, st)
     ϕe = StatefulLuxLayer{true}(l.ϕe, ps.ϕe, _getstate(st, :ϕe))
-    ϕv = StatefulLuxLayer{true}(l.ϕv, ps.ϕv, _getstate(st, :ϕv))    
+    ϕv = StatefulLuxLayer{true}(l.ϕv, ps.ϕv, _getstate(st, :ϕv))
     m = (; ϕe, ϕv, aggr=l.aggr)
-    return GNNlib.megnet_conv(m, g, x, e), st
+    y = GNNlib.megnet_conv(m, g, x, e)
+    stnew = (; ϕe = _getstate(ϕe), ϕv = _getstate(ϕv))
+    return y, stnew
 end
 
 

GNNLux/test/layers/conv.jl

@@ -107,6 +107,17 @@
 
         @test size(x_new) == (out_dims, g.num_nodes)
         @test size(e_new) == (out_dims, g.num_edges)
+
+        ϕe = Chain(Dense(3in_dims => out_dims), BatchNorm(out_dims))
+        ϕv = Chain(Dense(in_dims + out_dims => out_dims), Dense(out_dims => out_dims))
+        l = MEGNetConv(in_dims, out_dims, ϕe, ϕv)
+        ps, st = LuxCore.setup(rng, l)
+        _, st_new = l(g, x, e, ps, st)
+
+        @test typeof(st_new) == typeof(st)
+        @test st_new.ϕe.layer_2.running_mean != st.ϕe.layer_2.running_mean
+        _, st_newer = l(g, x, e, ps, st_new)
+        @test typeof(st_newer) == typeof(st)
     end
 
     @testset "NNConv" begin

docs/CONTRIBUTION_ROADMAP.md

@@ -0,0 +1,100 @@
+# 30 / 60 / 90-Day Contribution Roadmap
+
+This plan assumes about 6–8 focused hours each week alongside study/work. Adjust the pace to your schedule. The goal is to learn by making small, reviewed changes; GSoC 2027 opportunities and mentor availability are not yet known.
+
+## Days 1–30: Build a working mental model
+
+### Week 1 — Julia and local workflow
+
+- Read `AGENTS.md`, `README.md`, and `DEVELOPMENT.md`.
+- Use `scripts/julia.ps1` to inspect `GNNGraph`, `DataStore`, and the Julia package environment.
+- Learn multiple dispatch, parametric structs, keyword arguments, broadcasting, named tuples, and Julia package extensions from the Julia manual.
+- Run one focused `@testitem`; then run the `GNNGraphs` CPU suite when data/network requirements are available.
+
+**Deliverable:** personal notes explaining `features × nodes`, COO source/target arrays, and how a test item is launched.
+
+### Week 2 — Graph representation and operations
+
+- Trace `GNNGraph` construction from `GNNGraphs/src/gnngraph.jl` into `to_coo`, `to_sparse`, `to_dense`, and `normalize_graphdata`.
+- Read `query.jl` for `edge_index`, `degree`, and adjacency semantics.
+- Read `transform.jl` batch/unbatch, self-loop, remove-edge, and `getgraph` functions.
+- Run the matching graph tests and inspect any `broken` annotations before interpreting results.
+
+**Deliverable:** a one-page graph API map with one homogeneous and one heterogeneous example.
+
+### Week 3 — Message passing and one layer
+
+- Read `GNNlib/src/msgpass.jl`; step through `apply_edges` → message function → `aggregate_neighbors`.
+- Compare a generic message closure with `copy_xj` and identify when GNNlib can use a fast path.
+- Follow GCN from Flux and Lux frontend calls to `GNNlib.gcn_conv`, then to GNNGraphs `degree`.
+- Compare the source code against layer docstrings and hand-calculate a tiny graph result.
+
+**Deliverable:** a trace diagram for GCN and a small manual reference calculation.
+
+### Week 4 — First small contribution
+
+- Read an open issue page and check for current branches/PRs before choosing work.
+- Recommended starting scope: a missing expected-value test, a documentation correction, or a reproducible error-message improvement.
+- Ask maintainers on the issue if the task is still wanted; do not duplicate an existing PR.
+- Write the regression test first, make the smallest change, run the package suite, and open a draft PR if appropriate.
+
+**Deliverable:** one small reviewable PR or a maintainer-confirmed test plan.
+
+## Days 31–60: Own one subsystem
+
+### Weeks 5–6 — Learn the test and AD model
+
+- Study `test_module.jl` for finite differences, Zygote, Mooncake and device selection.
+- Pick one layer and inspect primal output tests plus gradient tests. Add a hand-computed reference to your local learning branch if a formula is otherwise only shape/gradient checked.
+- Learn how `@testitem` setup modules and `GNN_TEST_*` flags select CPU/GPU execution.
+- Read relevant Mooncake/Enzyme extensions and the changelog entries for prior AD fixes.
+
+**Deliverable:** a backend/layer/graph-format test matrix for one operation.
+
+### Weeks 7–8 — Reproduce and scope a real issue
+
+- Check current issue #710 or #640 status and linked PRs; verify the versions/OS/devices in the issue before running it.
+- Build a minimal script outside production source. Capture Julia/package versions and stack trace.
+- Split graph data, message passing, frontend and backend responsibilities. Use `@which`, `methods`, and `@code_warntype` as appropriate; do not treat type instability output as proof of a bug by itself.
+- Present reproduction and test plan to maintainers before proposing a backend-wide change.
+
+**Deliverable:** a minimal reproducible example, root-cause evidence, and scoped implementation plan.
+
+## Days 61–90: Deliver a medium contribution and prepare a proposal
+
+### Weeks 9–10 — Implement with regression coverage
+
+- Implement in the lowest package that owns the behavior: graph semantics in GNNGraphs, shared math in GNNlib, training/state wrapper in Flux/Lux frontend, backend-specialized behavior in an extension.
+- Keep CPU behavior and optional backend behavior tested separately.
+- Add exact outputs/gradients or reference results, including isolated node, empty-edge, directed graph and multiple representation cases when relevant.
+- Benchmark only if performance is the stated goal; measure warm and steady-state behavior, allocations and device synchronization.
+
+**Deliverable:** focused PR with reproduction, tests, review notes and accurate limitations.
+
+### Weeks 11–12 — GSoC exploration
+
+- Identify one problem theme and ask a maintainer whether it fits the project's likely priorities; do not assume a 2027 slot exists.
+- Write a short design note: user impact, package ownership, API surface, baseline, test plan, milestones, risks, fallback and mentor touchpoints.
+- Read current GSoC/Julia organization dates and application guidance once published.
+- Contribute a second small PR or review an existing PR to demonstrate collaboration and follow-through.
+
+**Deliverable:** proposal draft grounded in a merged/under-review contribution and current maintainer feedback.
+
+## 90-day success criteria
+
+- Can trace a graph operation from public API through the owning package and tests.
+- Can explain the graph source/target orientation and feature axis convention.
+- Has run a package CPU suite and at least one focused test item, recording any environment failures honestly.
+- Has one or more small PRs or a clear maintainer-approved implementation plan.
+- Can distinguish a confirmed bug from a report, a TODO, a backend gap, and a performance hypothesis.
+- Has a GSoC project outline with scoped milestones, but no assumption of acceptance or project availability.
+
+## Weekly routine
+
+1. Sync your fork with upstream `master` and create a focused branch.
+2. Read the owning package's source, matching test, docs and recent changelog/history.
+3. Reproduce before changing code.
+4. Make one minimal change and run the smallest relevant test first.
+5. Run the full package CPU suite and any available optional backend suite.
+6. Review the diff for accidental generated files, summarize evidence, then ask maintainers for review.
+