[DHAKED440]

Fix #710

Summary

Fix Reactant compilation failures in GCNConv graph normalization and message passing.

Reactant traced scalar types are now mapped to their underlying numeric types before graph-side degree and adjacency allocations, ensuring compatibility with Reactant compilation.

Reactant support is provided through an optional GNNGraphs extension, keeping Reactant optional for normal CPU users.

Changes

• Fix degree calculation and message-passing allocations for Reactant traced scalar types.
• Add optional Reactant integration through a GNNGraphs extension.
• Add integration tests covering degree compilation, GCNConv compilation and execution, output shape, and numerical agreement with the CPU result.

Tests

All executed test suites passed:

• Optional Reactant integration tests: Passed
• GNNGraphs: 1,352 passed, 8 broken
• GNNlib: 86 passed
• GNNLux: 177 passed, 2 broken
• GraphNeuralNetworks: 928 passed

The compiled GCNConv output has the expected shape (2, 6) and matches the CPU result within the test tolerance.

[simsurace]

Thanks for the PR. Since in this approach the graph is a compile-time constant, it may give silently wrong results when the compiled function is applied to a different graph.

Reproducer: a compiled model silently reuses the graph it was compiled with

using GNNLux, Lux, Reactant, Random

model = GCNConv(3 => 2)
ps, st = Lux.setup(Xoshiro(0), model)
x = randn(Xoshiro(1), Float32, 3, 6)
g1 = GNNGraph(([1, 2, 3, 4, 5, 6], [2, 3, 1, 5, 6, 4]); num_nodes = 6)
g2 = GNNGraph(([1, 1, 1, 2, 2, 3], [4, 5, 6, 5, 6, 6]); num_nodes = 6)

f(g, x, ps, st) = first(model(g, x, ps, st))
ps_d, st_d, x_d = (ps, st, x) |> reactant_device()
compiled = @compile f(g1, x_d, ps_d, st_d)   # graph stays on the host

y = Array(compiled(g2, x_d, ps_d, st_d))
y ≈ f(g2, x, ps, st)   # false
y ≈ f(g1, x, ps, st)   # true: g1 was baked into the compiled program

The host graph's edge indices become constants in the compiled program, so calling it with another graph of the same type gives no error and returns the result for the original graph. This affects mini-batched training and any compiled model reused across graphs.

#730 is more general in that respect, i.e. it goes further by ensuring that the graph can be traced, which is needed for the tutorials to fully work.

With #730: device graphs are inputs, so a compiled model works across graphs

We also move the graphs to the device; this would fail on this branch:

using GNNLux, Lux, Reactant, Random

model = GCNConv(3 => 2)
ps, st = Lux.setup(Xoshiro(0), model)
x = randn(Xoshiro(1), Float32, 3, 6)
g1 = GNNGraph(([1, 2, 3, 4, 5, 6], [2, 3, 1, 5, 6, 4]); num_nodes = 6)
g2 = GNNGraph(([1, 1, 1, 2, 2, 3], [4, 5, 6, 5, 6, 6]); num_nodes = 6)

f(g, x, ps, st) = first(model(g, x, ps, st))
ps_d, st_d, x_d, g1_d, g2_d = (ps, st, x, g1, g2) |> reactant_device()
compiled = @compile f(g1_d, x_d, ps_d, st_d)   # graph moved to the device

Array(compiled(g1_d, x_d, ps_d, st_d)) ≈ f(g1, x, ps, st)   # true
Array(compiled(g2_d, x_d, ps_d, st_d)) ≈ f(g2, x, ps, st)   # true

The graph's edge indices are inputs to the compiled program, so the same compiled model gives correct results for any graph with the same numbers of nodes and edges. Graphs of a different size need a new @compile.

That said, even #730 does not correctly deal with different batch sizes yet.

Note that the Reactant tests here are in a separate environment that is not part of CI.

[DHAKED440]

Thanks for catching this. You’re right: my current test compiles with a host graph, so its edge indices can be baked into the compiled program. Reusing it with another graph could silently produce the result for the original topology. I’ll investigate the device-graph approach and compare with #730 so the graph is traced as a runtime input, then add a test that reuses the compiled function across same-shaped graphs. I’ll also account for the need to recompile when graph shapes change.

[CarloLucibello]

Thanks, nice work getting this through Reactant! I checked the branch out and ran it locally (Julia 1.13.3, Reactant 0.2.307).

Verified

• The three Reactant testsets pass.
• Beyond the PR's tests, these also compile and match CPU: GCNConv on a weighted graph, with use_edge_weight=true, and with add_self_loops=false, plus GraphConv and SAGEConv.
• The propagate overrides mirror the CUDA extension, including the w_mul_xj weights fix.

Requested change

1. Relaxing COO_T, AVecI and _scatter to AbstractVector removes integer-index validation for all users, not just Reactant ones. The comment "Constructors still validate ordinary graph indices" does not hold:

   • GNNGraph([1.0, 2.0], [2.0, 1.0]; num_nodes=2) now builds a graph (so does [1.5, 2.0]), and fails later in degree / gather with invalid index: 1.0 of type Float64. On master this is a clear MethodError at construction.
   • String indices now give isless(::Int64, ::String) instead of a MethodError.

   Could you keep the relaxed aliases but validate the element type, defaulting to Integer with the extension allowing traced integers? A rough sketch (untested on the traced paths, so please check it against the Reactant tests):

   # GNNGraphs, called from to_coo(::COO_T) for s and t
   _check_index_eltype(::Type{T}) where {T} =
       T <: Integer || throw(ArgumentError("graph indices must be integers, got eltype $T"))
   # GNNGraphsReactantExt
   GNNGraphs._check_index_eltype(::Type{<:Reactant.TracedRNumber{<:Integer}}) = true

Suggestions
2. The new tests live in the opt-in GNNLux/test/reactant environment, which no workflow runs, so they can regress silently (this is also why codecov/patch is red). Consider a CI job for it. It needs Julia ≥ 1.11 because of [sources], while that Project.toml declares julia = "1.10".
3. GNNGraphsReactantExt relies on Reactant.unwrapped_eltype. It works on 0.2.307, but compat is Reactant = "0.2". Matching TracedRNumber{T} and returning T directly would not depend on a Reactant internal.
4. Nit: the _scatter added to the using .GNNGraphs: list in GNNlib/src/GNNlib.jl is unused; msgpass.jl calls GNNGraphs._scatter.

🧔‍♂️🤖 Assisted by AI

[DHAKED440]

Fixes #710

Summary

Fix Reactant compilation for GCNConv while keeping graph topology available as a runtime input to compiled models.

Root cause

When a host-side graph is used during compilation, its edge indices can be embedded as constants. Reusing the compiled function with another graph of the same shape may then silently use the original topology.

Changes

• Allow Reactant traced integer types for COO graph indices and graph indicators while preserving integer-index validation for ordinary graphs.
• Use gather/scatter message passing for Reactant traced graphs through an optional GNNlib extension.
• Use the wrapped numeric type directly in the GNNGraphs Reactant extension.
• Add regression tests for reusing a compiled model with a different, same-shaped graph.
• Add constructor tests rejecting floating-point and string graph indices and graph indicators.
• Add a Julia 1.11 CI job for the optional Reactant integration tests.
• Remove an unused _scatter import.

Tests

• git diff --check: passed.
• Project TOML parsing: passed.
• Julia package tests: not completed locally because Pkg could not access its depot in this environment.
• GitHub CI, including the Reactant job: running.