Fix Reactant compilation for GCNConv

.github/workflows/test_GNNLux_Reactant.yml

@@ -0,0 +1,38 @@
+name: GNNLux Reactant
+
+on:
+  pull_request:
+    branches:
+      - master
+  push:
+    branches:
+      - master
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/master' }}
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    name: Julia 1.11 - Reactant CPU
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: julia-actions/setup-julia@v3
+        with:
+          version: '1.11'
+      - uses: julia-actions/cache@v3
+      - name: Install Reactant test dependencies
+        run: julia --project=GNNLux/test/reactant -e 'using Pkg; Pkg.instantiate()'
+      - name: Run optional Reactant integration tests
+        run: julia --code-coverage=all --project=GNNLux/test/reactant GNNLux/test/reactant/runtests.jl
+      - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: GNNLux/src,GNNlib/src,GNNlib/ext,GNNGraphs/src,GNNGraphs/ext
+      - uses: codecov/codecov-action@v7
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: lcov.info

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 — Unreleased (towards 1.7.0)
+
+**Fixed**
+- Reactant traced scalar types are now converted to their underlying numeric type before graph-side degree and COO adjacency allocation. COO graph indices can also remain traced device inputs during compiled GNNLux model calls. Reactant remains an optional GNNGraphs integration ([#710]).
+
 ## GNNGraphs.jl 1.6.0 - 2026-11-04
 
 **Added**
@@ -27,6 +32,11 @@ Entries link to the pull request that introduced them.
 - `GNNChain` now returns its state with the same structure it receives, `(; layers = (; layer_1, ...))`, as the Lux layer interface requires. It previously returned the per-layer states without the `layers` wrapper, so feeding the returned state back into the chain (as `Lux.Training.single_train_step!` does) failed with a `FieldError`. This is breaking for code that worked around the bug by re-wrapping the state as `(; layers = st)`, as the docs examples did: drop the wrapper.
 - `NNConv` now returns its state as `(; nn = ...)`, the same structure it receives. It previously returned the inner network's state without the `nn` wrapper, so passing the returned state back into the layer (e.g. on the second step of `Lux.Training.single_train_step!`) failed with a `FieldError` ([#728]).
 
+## GNNlib.jl — Unreleased (towards 1.5.0)
+
+**Fixed**
+- Reactant compiled message passing now uses gather/scatter with traced COO indices instead of building a sparse adjacency matrix, allowing one compiled GCNConv to run on different graphs with the same shapes ([#710]).
+
 ## GNNlib.jl 1.4.1 - 2026-11-04
 
 **Fixed**
@@ -251,4 +261,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
 [#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722
 [#718]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/718
 [#728]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/728
+[#710]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/710
 [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.6.0"
+version = "1.7.0-DEV"
 authors = ["Carlo Lucibello and contributors"]
 
 [workspace]
@@ -25,12 +25,14 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 
 [extensions]
 GNNGraphsCUDAExt = "CUDA"
 GNNGraphsEnzymeCoreExt = "EnzymeCore"
 GNNGraphsMooncakeExt = "Mooncake"
+GNNGraphsReactantExt = "Reactant"
 GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
 
 [compat]
@@ -46,6 +48,7 @@ MLUtils = "0.4"
 Mooncake = "0.5.24"
 NNlib = "0.9"
 NearestNeighbors = "0.4"
+Reactant = "0.2"
 Random = "1"
 SimpleWeightedGraphs = "1.4.0"
 SparseArrays = "1"

GNNGraphs/ext/GNNGraphsReactantExt.jl

@@ -0,0 +1,10 @@
+module GNNGraphsReactantExt
+
+using GNNGraphs
+using Reactant: TracedRNumber
+
+GNNGraphs._check_index_eltype(::Type{<:TracedRNumber{T}}) where {T <: Integer} = true
+
+GNNGraphs._graph_scalar_eltype(::Type{TracedRNumber{T}}) where {T} = T
+
+end

GNNGraphs/src/abstracttypes.jl

@@ -1,10 +1,16 @@
 
-const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V <: Union{Nothing, AbstractVector}}
+# COO index arrays may have non-integer wrapper eltypes while traced by a compiler
+# such as Reactant. `_check_index_eltype` validates that the wrapped type is integral.
+const COO_T = Tuple{T, T, V} where {T <: AbstractVector, V <: Union{Nothing, AbstractVector}}
 const ADJLIST_T = AbstractVector{T} where {T <: AbstractVector{<:Integer}}
 const ADJMAT_T = AbstractMatrix
 const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T
 
-const AVecI = AbstractVector{<:Integer}
+# Graph indicator arrays may likewise be traced arrays whose eltype wraps integers.
+const AVecI = AbstractVector
+
+_check_index_eltype(::Type{T}) where {T} =
+    T <: Integer || throw(ArgumentError("graph indices must be integers, got eltype $T"))
 
 # All concrete graph types should be subtypes of AbstractGNNGraph{T}.
 # GNNGraph and GNNHeteroGraph are the two concrete types.

GNNGraphs/src/convert.jl

@@ -30,6 +30,9 @@ function to_coo(coo::COO_T; dir = :out, num_nodes = nothing, weighted = true,
                 hetero = false)
     s, t, val = coo
 
+    _check_index_eltype(eltype(s))
+    _check_index_eltype(eltype(t))
+
     if isnothing(num_nodes)
         ns = maximum(s)
         nt = maximum(t)
@@ -146,6 +149,7 @@ function to_dense(A::ADJMAT_T, T = nothing; dir = :out, num_nodes = nothing,
                   weighted = true)
     @assert dir ∈ [:out, :in]
     T = T === nothing ? eltype(A) : T
+    T = _graph_scalar_eltype(T)
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
     # @assert all(x -> (x == 1) || (x == 0), A)
@@ -190,6 +194,7 @@ function to_dense(coo::COO_T, T = nothing; dir = :out, num_nodes = nothing, weig
     if T === nothing
         T = isnothing(val) ? eltype(s) : eltype(val)
     end
+    T = _graph_scalar_eltype(T)
     if val === nothing || !weighted
         val = ones_like(s, T)
     end
@@ -216,6 +221,7 @@ function to_sparse(A::ADJMAT_T, T = nothing; dir = :out, num_nodes = nothing,
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
     T = T === nothing ? eltype(A) : T
+    T = _graph_scalar_eltype(T)
     num_edges = A isa AbstractSparseMatrix ? nnz(A) : count(!=(0), A)
     if dir == :in
         A = A'
@@ -242,6 +248,7 @@ function to_sparse(coo::COO_T, T = nothing; dir = :out, num_nodes = nothing,
                    weighted = true)
     s, t, eweight = coo
     T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T
+    T = _graph_scalar_eltype(T)
 
     if eweight === nothing || !weighted
         eweight = fill!(similar(s, T), 1)

GNNGraphs/src/gatherscatter.jl

@@ -11,7 +11,7 @@ _scatter(aggr, src::Dict, idx, n) = Dict([k => _scatter(aggr, v, idx, n) for (k,
 
 function _scatter(aggr,
                   src::AbstractArray,
-                  idx::AbstractVector{<:Integer},
+                  idx::AbstractVector,
                   n::Integer)
     dstsize = (size(src)[1:(end - 1)]..., n)
     return NNlib.scatter(aggr, src, idx; dstsize)

GNNGraphs/src/gnngraph.jl

@@ -125,6 +125,9 @@ function GNNGraph(graph::T,
                   ndata::DataStore,
                   edata::DataStore,
                   gdata::DataStore) where {T <: Union{COO_T, ADJMAT_T}}
+    if !isnothing(graph_indicator)
+        _check_index_eltype(eltype(graph_indicator))
+    end
     return GNNGraph{T}(graph, num_nodes, num_edges, num_graphs,
                        graph_indicator, ndata, edata, gdata, false)
 end
@@ -154,6 +157,9 @@ function GNNGraph(data::D;
         graph, num_nodes, num_edges = to_sparse(data; num_nodes, dir)
     end
 
+    if !isnothing(graph_indicator)
+        _check_index_eltype(eltype(graph_indicator))
+    end
     num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1
 
     ndata = normalize_graphdata(ndata, default_name = :x, n = num_nodes)

GNNGraphs/src/gnnheterograph/gnnheterograph.jl

@@ -117,6 +117,10 @@ function GNNHeteroGraph(data::EDict;
     @assert dir ∈ [:in, :out]
     @assert graph_type==:coo "only :coo graph_type is supported for now"
 
+    if !isnothing(graph_indicator)
+        foreach(gi -> _check_index_eltype(eltype(gi)), values(graph_indicator))
+    end
+
     if num_nodes !== nothing
         num_nodes = Dict(num_nodes)
     end

GNNGraphs/src/query.jl

@@ -350,6 +350,7 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T::TT = nothing; dir = :out,
         else 
             T
         end
+    T = _graph_scalar_eltype(T)
     return _degree((s, t), T, dir, ew, g.num_nodes)
 end
 
@@ -371,10 +372,15 @@ function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T::TT = nothing; dir = :out,
             T = Nt
         end
     end
+    T = _graph_scalar_eltype(T)
     A = adjacency_matrix(g)
     return _degree(A, T, dir, edge_weight, g.num_nodes)
 end
 
+# Backends can specialize this hook when their scalar wrapper type should not be
+# used as the element type of ordinary graph-side allocations.
+_graph_scalar_eltype(::Type{T}) where {T} = T
+
 function _degree((s, t)::Tuple, T::Type, dir::Symbol, edge_weight::Nothing, num_nodes::Int)
     _degree((s, t), T, dir, ones_like(s, T), num_nodes)
 end

GNNGraphs/test/gnngraph.jl

@@ -1,6 +1,20 @@
 # TODO test that the graph type is preserved
 # when constructing a GNNGraph from another
 
+@testitem "Constructor: COO indices must be integers" setup=[GraphsTestModule] begin
+    @test GNNGraph(([1, 2], [2, 1]); num_nodes = 2).num_edges == 2
+    @test_throws ArgumentError GNNGraph(([1.0, 2.0], [2.0, 1.0]); num_nodes = 2)
+    @test_throws ArgumentError GNNGraph(([1.5, 2.0], [2.0, 1.0]); num_nodes = 2)
+    @test_throws ArgumentError GNNGraph((["1", "2"], ["2", "1"]); num_nodes = 2)
+
+    @test_throws ArgumentError GNNGraph(
+        ([1, 2], [2, 1]); num_nodes = 2, graph_indicator = [1.0, 1.0]
+    )
+    @test_throws ArgumentError GNNGraph(
+        ([1, 2], [2, 1]); num_nodes = 2, graph_indicator = ["1", "1"]
+    )
+end
+
 @testitem "Constructor: adjacency matrix" setup=[GraphsTestModule] begin
     using .GraphsTestModule
     for GRAPH_T in GRAPH_TYPES

GNNLux/test/reactant/Project.toml

@@ -0,0 +1,17 @@
+[deps]
+GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
+GNNLux = "e8545f4d-a905-48ac-a8c4-ca114b98986d"
+GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+GNNGraphs = {path = "../../../GNNGraphs"}
+GNNLux = {path = "../.."}
+GNNlib = {path = "../../../GNNlib"}
+
+[compat]
+Reactant = "0.2"
+julia = "1.11"

GNNLux/test/reactant/runtests.jl

@@ -0,0 +1,76 @@
+# Optional integration tests; the regular GNNLux test project does not depend on
+# Reactant. From the repository root, instantiate and run with:
+#   julia --project=GNNLux/test/reactant -e 'using Pkg; Pkg.instantiate()'
+#   julia --project=GNNLux/test/reactant GNNLux/test/reactant/runtests.jl
+
+using GNNGraphs
+using GNNLux
+using Lux
+using Random
+using Reactant
+using Test
+
+degree_for_features(g, x) = GNNGraphs.degree(g, eltype(x); dir = :in, edge_weight = false)
+
+function test_graph()
+    return GNNGraph(
+        ([1, 2, 3, 4, 5, 6], [2, 3, 1, 5, 6, 4]);
+        num_nodes = 6,
+    )
+end
+
+@testset "Reactant degree uses feature numeric type" begin
+    g = test_graph()
+    x = randn(Xoshiro(42), Float32, 3, 6)
+    x_d = x |> reactant_device()
+
+    compiled_degree = @compile degree_for_features(g, x_d)
+    degree_d = compiled_degree(g, x_d)
+
+    @test Array(degree_d) == ones(Float32, 6)
+end
+
+@testset "Reactant GCNConv compiles, executes, and matches CPU" begin
+    rng = Xoshiro(42)
+    model = GCNConv(3 => 2, relu)
+    ps, st = Lux.setup(rng, model)
+    g = test_graph()
+    x = randn(rng, Float32, 3, 6)
+    y_cpu, _ = model(g, x, ps, st)
+
+    dev = reactant_device()
+    ps_d, st_d = (ps, st) |> dev
+    x_d = x |> dev
+    compiled = @compile model(g, x_d, ps_d, st_d)
+    y_d, _ = compiled(g, x_d, ps_d, st_d)
+
+    @test size(y_d) == (2, 6)
+    @test Array(y_d) ≈ y_cpu rtol = 1.0f-4 atol = 1.0f-5
+end
+
+@testset "Reactant GCNConv keeps device graph indices dynamic" begin
+    rng = Xoshiro(42)
+    model = GCNConv(3 => 2, relu)
+    ps, st = Lux.setup(rng, model)
+    g1 = test_graph()
+    g2 = GNNGraph(
+        ([1, 1, 1, 2, 2, 3], [4, 5, 6, 5, 6, 6]);
+        num_nodes = 6,
+    )
+    x = randn(rng, Float32, 3, 6)
+    f(g, x, ps, st) = first(model(g, x, ps, st))
+
+    dev = reactant_device()
+    ps_d, st_d, x_d, g1_d, g2_d = (ps, st, x, g1, g2) |> dev
+    compiled = @compile f(g1_d, x_d, ps_d, st_d)
+
+    y1_cpu = f(g1, x, ps, st)
+    y2_cpu = f(g2, x, ps, st)
+    y1_compiled = Array(compiled(g1_d, x_d, ps_d, st_d))
+    y2_compiled = Array(compiled(g2_d, x_d, ps_d, st_d))
+
+    @test size(y1_compiled) == (2, 6)
+    @test y1_compiled ≈ y1_cpu rtol = 1.0f-4 atol = 1.0f-5
+    @test size(y2_compiled) == (2, 6)
+    @test y2_compiled ≈ y2_cpu rtol = 1.0f-4 atol = 1.0f-5
+end

GNNlib/Project.toml

@@ -1,7 +1,7 @@
 name = "GNNlib"
 uuid = "a6a84749-d869-43f8-aacc-be26a1996e48"
 authors = ["Carlo Lucibello and contributors"]
-version = "1.4.1"
+version = "1.5.0-DEV"
 
 [workspace]
 projects = ["test", "docs", "../GNNGraphs"]
@@ -21,24 +21,27 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]
 GNNlibAMDGPUExt = "AMDGPU"
 GNNlibCUDAExt = "CUDA"
 GNNlibMetalExt = "Metal"
 GNNlibMooncakeExt = "Mooncake"
+GNNlibReactantExt = "Reactant"
 
 [compat]
 AMDGPU = "1, 2.7"
 CUDA = "5, 6"
 ChainRulesCore = "1.24"
 DataStructures = "0.18, 0.19"
-GNNGraphs = "1.6"
+GNNGraphs = "1.7"
 LinearAlgebra = "1"
 Metal = "1.0"
 MLUtils = "0.4"
 Mooncake = "0.5"
 NNlib = "0.9"
 Random = "1"
+Reactant = "0.2"
 Statistics = "1"
 julia = "1.10"