From dbf40d3cd67f3bdcf31503d93e13815f880734f0 Mon Sep 17 00:00:00 2001 From: Rahul dhaked Date: Sat, 7 Nov 2026 22:45:06 +0530 Subject: [PATCH 1/4] Support Reactant scalar types in graph allocation --- CHANGELOG.md | 2 ++ GNNGraphs/Project.toml | 3 ++ GNNGraphs/ext/GNNGraphsReactantExt.jl | 9 +++++ GNNGraphs/src/convert.jl | 4 +++ GNNGraphs/src/query.jl | 6 ++++ GNNLux/test/reactant/Project.toml | 17 ++++++++++ GNNLux/test/reactant/runtests.jl | 49 +++++++++++++++++++++++++++ 7 files changed, 90 insertions(+) create mode 100644 GNNGraphs/ext/GNNGraphsReactantExt.jl create mode 100644 GNNLux/test/reactant/Project.toml create mode 100644 GNNLux/test/reactant/runtests.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 7611c27fe..3105348aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ Entries link to the pull request that introduced them. **Fixed** - `remove_nodes` now updates the `graph_indicator` of batched graphs. It was previously left unchanged, so calling `unbatch` or `getgraph` on the result threw a `BoundsError` ([#717]). +- Reactant traced scalar types are now converted to their underlying numeric type before graph-side degree and COO adjacency allocation. Reactant remains an optional GNNGraphs integration ([#710]). - Fixed `scaled_laplacian` (and so `ChebConv`) crashing with a segmentation fault on AMDGPU graphs: the eigensolver's start vector was created on the CPU, so CPU BLAS was handed GPU memory. It is now created on the same device as the graph ([#722]). - Added an `EnzymeCore` package extension marking `scaled_laplacian` inactive for Enzyme, matching its existing `@non_differentiable` ChainRules declaration. Enzyme previously differentiated the Krylov eigensolve inside it and failed; this unblocks `ChebConv`, `GConvGRUCell` and `GConvLSTMCell` ([#706]). - Enzyme can now differentiate through the adjacency-matrix → COO graph conversion: a new internal keyword-free helper `_to_coo_graph` avoids the union-typed keyword handling of the `GNNGraph(g; graph_type)` constructor and of `to_coo`, which Enzyme's type analysis cannot compile. Together with upstream fixes in Enzyme ≥ 0.13.197 this removes the `EnzymeInternalError` crash on `:dense`/`:sparse` graphs ([#703]). @@ -251,4 +252,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 diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml index 29c2d2272..bcd94346c 100644 --- a/GNNGraphs/Project.toml +++ b/GNNGraphs/Project.toml @@ -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" diff --git a/GNNGraphs/ext/GNNGraphsReactantExt.jl b/GNNGraphs/ext/GNNGraphsReactantExt.jl new file mode 100644 index 000000000..50cab4de6 --- /dev/null +++ b/GNNGraphs/ext/GNNGraphsReactantExt.jl @@ -0,0 +1,9 @@ +module GNNGraphsReactantExt + +using GNNGraphs +using Reactant + +GNNGraphs._graph_scalar_eltype(::Type{Reactant.TracedRNumber{T}}) where {T} = + Reactant.unwrapped_eltype(Reactant.TracedRNumber{T}) + +end diff --git a/GNNGraphs/src/convert.jl b/GNNGraphs/src/convert.jl index f5675ffff..efb97a1f9 100644 --- a/GNNGraphs/src/convert.jl +++ b/GNNGraphs/src/convert.jl @@ -146,6 +146,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 +191,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 +218,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 +245,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) diff --git a/GNNGraphs/src/query.jl b/GNNGraphs/src/query.jl index aaad2d634..db8b4ac98 100644 --- a/GNNGraphs/src/query.jl +++ b/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 diff --git a/GNNLux/test/reactant/Project.toml b/GNNLux/test/reactant/Project.toml new file mode 100644 index 000000000..5c2e66896 --- /dev/null +++ b/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.10" diff --git a/GNNLux/test/reactant/runtests.jl b/GNNLux/test/reactant/runtests.jl new file mode 100644 index 000000000..0f198202a --- /dev/null +++ b/GNNLux/test/reactant/runtests.jl @@ -0,0 +1,49 @@ +# 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 From 00dbc4c1ac873b13fdd3776ff6adefdf23eeb61b Mon Sep 17 00:00:00 2001 From: Rahul dhaked Date: Sat, 7 Nov 2026 22:45:36 +0530 Subject: [PATCH 2/4] Keep Reactant graph indices dynamic in GCNConv --- CHANGELOG.md | 11 ++++++++++- GNNGraphs/Project.toml | 2 +- GNNGraphs/src/abstracttypes.jl | 7 +++++-- GNNGraphs/src/gatherscatter.jl | 2 +- GNNLux/test/reactant/runtests.jl | 27 +++++++++++++++++++++++++++ GNNlib/Project.toml | 7 +++++-- GNNlib/ext/GNNlibReactantExt.jl | 25 +++++++++++++++++++++++++ GNNlib/src/GNNlib.jl | 3 +-- 8 files changed, 75 insertions(+), 9 deletions(-) create mode 100644 GNNlib/ext/GNNlibReactantExt.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 3105348aa..5c1704b9b 100644 --- a/CHANGELOG.md +++ b/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** @@ -17,7 +22,6 @@ Entries link to the pull request that introduced them. **Fixed** - `remove_nodes` now updates the `graph_indicator` of batched graphs. It was previously left unchanged, so calling `unbatch` or `getgraph` on the result threw a `BoundsError` ([#717]). -- Reactant traced scalar types are now converted to their underlying numeric type before graph-side degree and COO adjacency allocation. Reactant remains an optional GNNGraphs integration ([#710]). - Fixed `scaled_laplacian` (and so `ChebConv`) crashing with a segmentation fault on AMDGPU graphs: the eigensolver's start vector was created on the CPU, so CPU BLAS was handed GPU memory. It is now created on the same device as the graph ([#722]). - Added an `EnzymeCore` package extension marking `scaled_laplacian` inactive for Enzyme, matching its existing `@non_differentiable` ChainRules declaration. Enzyme previously differentiated the Krylov eigensolve inside it and failed; this unblocks `ChebConv`, `GConvGRUCell` and `GConvLSTMCell` ([#706]). - Enzyme can now differentiate through the adjacency-matrix → COO graph conversion: a new internal keyword-free helper `_to_coo_graph` avoids the union-typed keyword handling of the `GNNGraph(g; graph_type)` constructor and of `to_coo`, which Enzyme's type analysis cannot compile. Together with upstream fixes in Enzyme ≥ 0.13.197 this removes the `EnzymeInternalError` crash on `:dense`/`:sparse` graphs ([#703]). @@ -28,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** diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml index bcd94346c..3347d8833 100644 --- a/GNNGraphs/Project.toml +++ b/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] diff --git a/GNNGraphs/src/abstracttypes.jl b/GNNGraphs/src/abstracttypes.jl index 73146160f..6097a9a20 100644 --- a/GNNGraphs/src/abstracttypes.jl +++ b/GNNGraphs/src/abstracttypes.jl @@ -1,10 +1,13 @@ -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. Constructors still validate ordinary graph indices. +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 # All concrete graph types should be subtypes of AbstractGNNGraph{T}. # GNNGraph and GNNHeteroGraph are the two concrete types. diff --git a/GNNGraphs/src/gatherscatter.jl b/GNNGraphs/src/gatherscatter.jl index e897399ed..05aff8295 100644 --- a/GNNGraphs/src/gatherscatter.jl +++ b/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) diff --git a/GNNLux/test/reactant/runtests.jl b/GNNLux/test/reactant/runtests.jl index 0f198202a..c595a053d 100644 --- a/GNNLux/test/reactant/runtests.jl +++ b/GNNLux/test/reactant/runtests.jl @@ -47,3 +47,30 @@ end @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 diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml index a0dd0597e..d2a979fba 100644 --- a/GNNlib/Project.toml +++ b/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" diff --git a/GNNlib/ext/GNNlibReactantExt.jl b/GNNlib/ext/GNNlibReactantExt.jl new file mode 100644 index 000000000..db27c3433 --- /dev/null +++ b/GNNlib/ext/GNNlibReactantExt.jl @@ -0,0 +1,25 @@ +module GNNlibReactantExt + +using Reactant: AnyTracedRMatrix +using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj +using GNNGraphs: GNNGraph, COO_T, get_edge_weight + +# The generic fast paths construct a sparse adjacency matrix, which Reactant +# cannot trace when its indices are dynamic. Gather/scatter keeps graph indices +# as inputs to the compiled program. +function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+), + xi, xj::AnyTracedRMatrix, e) + propagate((xi, xj, e) -> copy_xj(xi, xj, e), g, +, xi, xj, e) +end + +function GNNlib.propagate(::typeof(e_mul_xj), g::GNNGraph{<:COO_T}, ::typeof(+), + xi, xj::AnyTracedRMatrix, e::AbstractVector) + propagate((xi, xj, e) -> e_mul_xj(xi, xj, e), g, +, xi, xj, e) +end + +function GNNlib.propagate(::typeof(w_mul_xj), g::GNNGraph{<:COO_T}, ::typeof(+), + xi, xj::AnyTracedRMatrix, e::Nothing) + propagate((xi, xj, e) -> w_mul_xj(xi, xj, e), g, +, xi, xj, get_edge_weight(g)) +end + +end # module diff --git a/GNNlib/src/GNNlib.jl b/GNNlib/src/GNNlib.jl index 8fe51c9a8..45e4e2115 100644 --- a/GNNlib/src/GNNlib.jl +++ b/GNNlib/src/GNNlib.jl @@ -10,7 +10,7 @@ using ChainRulesCore: @non_differentiable, ignore_derivatives using GNNGraphs using .GNNGraphs: COO_T, ADJMAT_T, SPARSE_T, check_num_nodes, check_num_edges, - _to_coo_graph, + _to_coo_graph, _scatter, EType, NType # for heteroconvs include("utils.jl") @@ -77,4 +77,3 @@ export global_pool, # include("layers/heteroconv.jl") # no functional part at the moment end #module - \ No newline at end of file From 72cfc86796265fa6222dd2d81d7317d540908c0d Mon Sep 17 00:00:00 2001 From: Rahul dhaked Date: Sat, 7 Nov 2026 22:44:23 +0530 Subject: [PATCH 3/4] Validate traced graph indices without relaxing CPU checks --- .github/workflows/test_GNNLux_Reactant.yml | 31 +++++++++++++++++++ GNNGraphs/ext/GNNGraphsReactantExt.jl | 7 +++-- GNNGraphs/src/abstracttypes.jl | 5 ++- GNNGraphs/src/convert.jl | 3 ++ GNNGraphs/src/gnngraph.jl | 6 ++++ .../src/gnnheterograph/gnnheterograph.jl | 4 +++ GNNGraphs/test/gnngraph.jl | 14 +++++++++ GNNLux/test/reactant/Project.toml | 2 +- GNNlib/src/GNNlib.jl | 2 +- 9 files changed, 68 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/test_GNNLux_Reactant.yml diff --git a/.github/workflows/test_GNNLux_Reactant.yml b/.github/workflows/test_GNNLux_Reactant.yml new file mode 100644 index 000000000..244692a9d --- /dev/null +++ b/.github/workflows/test_GNNLux_Reactant.yml @@ -0,0 +1,31 @@ +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 --project=GNNLux/test/reactant GNNLux/test/reactant/runtests.jl diff --git a/GNNGraphs/ext/GNNGraphsReactantExt.jl b/GNNGraphs/ext/GNNGraphsReactantExt.jl index 50cab4de6..b74ced782 100644 --- a/GNNGraphs/ext/GNNGraphsReactantExt.jl +++ b/GNNGraphs/ext/GNNGraphsReactantExt.jl @@ -1,9 +1,10 @@ module GNNGraphsReactantExt using GNNGraphs -using Reactant +using Reactant: TracedRNumber -GNNGraphs._graph_scalar_eltype(::Type{Reactant.TracedRNumber{T}}) where {T} = - Reactant.unwrapped_eltype(Reactant.TracedRNumber{T}) +GNNGraphs._check_index_eltype(::Type{<:TracedRNumber{T}}) where {T <: Integer} = true + +GNNGraphs._graph_scalar_eltype(::Type{TracedRNumber{T}}) where {T} = T end diff --git a/GNNGraphs/src/abstracttypes.jl b/GNNGraphs/src/abstracttypes.jl index 6097a9a20..5153063a3 100644 --- a/GNNGraphs/src/abstracttypes.jl +++ b/GNNGraphs/src/abstracttypes.jl @@ -1,6 +1,6 @@ # COO index arrays may have non-integer wrapper eltypes while traced by a compiler -# such as Reactant. Constructors still validate ordinary graph indices. +# 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 @@ -9,6 +9,9 @@ const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T # 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. abstract type AbstractGNNGraph{T} <: AbstractGraph{Int} end diff --git a/GNNGraphs/src/convert.jl b/GNNGraphs/src/convert.jl index efb97a1f9..e78fec966 100644 --- a/GNNGraphs/src/convert.jl +++ b/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) diff --git a/GNNGraphs/src/gnngraph.jl b/GNNGraphs/src/gnngraph.jl index dccaf8d67..b24322e14 100644 --- a/GNNGraphs/src/gnngraph.jl +++ b/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) diff --git a/GNNGraphs/src/gnnheterograph/gnnheterograph.jl b/GNNGraphs/src/gnnheterograph/gnnheterograph.jl index 7c2b9ede5..9c00e0856 100644 --- a/GNNGraphs/src/gnnheterograph/gnnheterograph.jl +++ b/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 diff --git a/GNNGraphs/test/gnngraph.jl b/GNNGraphs/test/gnngraph.jl index ca9c8af02..534d5f938 100644 --- a/GNNGraphs/test/gnngraph.jl +++ b/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 diff --git a/GNNLux/test/reactant/Project.toml b/GNNLux/test/reactant/Project.toml index 5c2e66896..433c9c510 100644 --- a/GNNLux/test/reactant/Project.toml +++ b/GNNLux/test/reactant/Project.toml @@ -14,4 +14,4 @@ GNNlib = {path = "../../../GNNlib"} [compat] Reactant = "0.2" -julia = "1.10" +julia = "1.11" diff --git a/GNNlib/src/GNNlib.jl b/GNNlib/src/GNNlib.jl index 45e4e2115..0ae7892db 100644 --- a/GNNlib/src/GNNlib.jl +++ b/GNNlib/src/GNNlib.jl @@ -10,7 +10,7 @@ using ChainRulesCore: @non_differentiable, ignore_derivatives using GNNGraphs using .GNNGraphs: COO_T, ADJMAT_T, SPARSE_T, check_num_nodes, check_num_edges, - _to_coo_graph, _scatter, + _to_coo_graph, EType, NType # for heteroconvs include("utils.jl") From dc47f12458497f32b9b561e1a342a17df09964bc Mon Sep 17 00:00:00 2001 From: Rahul dhaked Date: Sun, 8 Nov 2026 11:00:09 +0530 Subject: [PATCH 4/4] Upload Reactant integration coverage to Codecov --- .github/workflows/test_GNNLux_Reactant.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_GNNLux_Reactant.yml b/.github/workflows/test_GNNLux_Reactant.yml index 244692a9d..b1a1b3608 100644 --- a/.github/workflows/test_GNNLux_Reactant.yml +++ b/.github/workflows/test_GNNLux_Reactant.yml @@ -28,4 +28,11 @@ jobs: - name: Install Reactant test dependencies run: julia --project=GNNLux/test/reactant -e 'using Pkg; Pkg.instantiate()' - name: Run optional Reactant integration tests - run: julia --project=GNNLux/test/reactant GNNLux/test/reactant/runtests.jl + 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