[DHAKED440]

Fixes #671

Summary

Fix the sample_neighbors and Zygote-through-add_edges failures found in CI.

Changes

• Handle nodes with no neighbors when sampling, including with replacement, while preserving normal sampling behavior.
• Keep reconstruction of empty graph metadata out of Zygote’s differentiated path while preserving gradients through edge weights.
• Add regression tests for empty and non-empty neighbor sampling, graph edge and weight preservation, and differentiation through add_edges.

Tests

• Focused sample_neighbors tests: 31 passed.
• Focused add_edges/Zygote tests: 16 passed.
• GNNGraphs CPU suite: 1,368 passed, 8 broken.
• GNNGraphs CI on Julia 1.10 and latest stable: passed.

Before and after

The CI failures involved sampling from an empty neighbor list and Zygote processing empty graph metadata. Both affected GNNGraphs CI checks now pass. The add_edges regression returns the expected gradient: Float32[1.0, 1.0].

[CarloLucibello]

Thanks for this! I checked the branch out and ran it locally (Julia 1.13.3).

Verified

• On master, differentiating through add_edges fails (map is not defined on sets: Zygote differentiates the keys(x1) comprehension in cat_features(::Dict, ::Dict)). On this branch the Differentiating across add_edges causes Zygote to segfault #671 reproducer returns Float32[1.0, 1.0]. It also works with non-empty edata, which fails on master too.
• sample_neighbors(...; replace=true) on zero-degree nodes: 3/300 random graphs failed on master, 0/300 here.

Suggestions

1. The datastore.jl change looks redundant. With only that file reverted, the Differentiating across add_edges causes Zygote to segfault #671 reproducer and my non-empty-edata gradient checks still passed, so the utils.jl change (iterating over a Vector of keys instead of keys(x1)) is what fixes it. The if CRC.ignore_derivatives() do … end also wraps an isempty check, which is not differentiable anyway. If it is needed for the Julia 1.12 segfault from the issue (I could only test 1.13), please keep it but simplify; this passed the same checks for me:

   if isempty(data1) && isempty(data2)
       return CRC.ignore_derivatives(() -> DataStore(n1 + n2, Dict{Symbol, Any}()))
   end

   Could you confirm on 1.12 which of the two is actually needed?
2. Please add a gradient test through add_edges with non-empty edata. The current test only covers empty metadata, but the utils.jl change is what makes the general case work.
3. The sample_neighbors fix is independent of Differentiating across add_edges causes Zygote to segfault #671 (the old test was flaky because rand_graph(10, 40) can have zero-degree nodes). Fine to keep here, but worth a line in the PR description.

🧔‍♂️🤖 Assisted by AI

[DHAKED440]

Summary

• Fix Zygote differentiation through add_edges when edge metadata is non-empty by iterating over a collected vector of metadata keys.
• Simplify handling of empty DataStore metadata.
• Handle zero-degree nodes in sample_neighbors(...; replace=true) without attempting to sample from an empty neighbor list.

Tests

• Add a gradient regression test for add_edges with non-empty edge metadata.
• Add coverage for sampling zero-degree nodes with replacement.

[CarloLucibello]

@DHAKED440 I didn't realize it before merging, but this PR contained a lot of AI-generated junk files. Please pay attention to what you commit

[DHAKED440]

Okay, I’ll review future PRs carefully and keep commits focused on the issue.