Support compiling GNNLux models with Reactant

CHANGELOG.md

@@ -10,6 +10,21 @@ 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)
+
+**Added**
+- Graphs can be compiled with Reactant.jl: a `GNNGraph` can now be moved to a Reactant device (e.g. `g |> reactant_device()`) and traced. The COO index vectors and the `graph_indicator` of a `GNNGraph` are no longer required to have an `Integer` eltype, so that traced arrays (eltype `TracedRNumber{Int}`) can be stored in them, and a new `Reactant` package extension makes `degree` work inside a trace ([#730]).
+- Added `pad_graph`, which pads a graph to given numbers of nodes, edges and graphs, putting the padding in extra graphs. Padding every mini-batch to the same size lets Reactant.jl compile a model or a training step once and reuse it for all batches ([#730]).
+- Added the `graph_sizes` field to `GNNGraph`, holding the number of nodes in each graph of a batch (`nothing` for a single graph). Its length is the number of graphs, so a model compiled with Reactant.jl now raises an error when called on a batch with a different number of graphs, instead of silently returning wrong results ([#730]).
+
+## GNNlib.jl — Unreleased (towards 1.5.0)
+
+**Added**
+- Added a `Reactant` package extension, so that layers relying on `propagate` (e.g. `GCNConv`) can be compiled with Reactant.jl and differentiated with Enzyme inside the compiled program. As on CUDA, the sparse adjacency-matrix fast paths of `propagate` are bypassed in favour of gather/scatter message passing. Requires GNNGraphs ≥ 1.7 ([#730]).
+
+**Fixed**
+- `reduce_nodes`, `reduce_edges`, `softmax_nodes`, `softmax_edges` and `softmax_edge_neighbors` now scatter into an output of explicit size (the number of graphs, or of nodes) instead of inferring it from the largest index. This makes `GlobalPool` and the attention layers compile with Reactant, and returns the correct number of columns for batched graphs whose trailing graphs have no nodes ([#730]).
+
 ## GNNGraphs.jl 1.6.0 - 2026-11-04
 
 **Added**
@@ -247,4 +262,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
 [#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706
 [#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722
 [#718]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/718
+[#730]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/730
 [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]
@@ -47,6 +49,7 @@ Mooncake = "0.5.24"
 NNlib = "0.9"
 NearestNeighbors = "0.4"
 Random = "1"
+Reactant = "0.2"
 SimpleWeightedGraphs = "1.4.0"
 SparseArrays = "1"
 Statistics = "1"

GNNGraphs/ext/GNNGraphsReactantExt.jl

@@ -0,0 +1,19 @@
+module GNNGraphsReactantExt
+
+using GNNGraphs
+using Reactant: TracedRNumber
+
+# Layers compute degrees with the feature eltype, which inside a Reactant trace is
+# `TracedRNumber{T}`. Allocating with that eltype from host index vectors would create a
+# `Vector{TracedRNumber{T}}`, so compute with the underlying number type instead.
+function GNNGraphs._degree((s, t)::Tuple, ::Type{TracedRNumber{T}}, dir::Symbol,
+                           edge_weight::Nothing, num_nodes::Int) where {T}
+    return GNNGraphs._degree((s, t), T, dir, edge_weight, num_nodes)
+end
+
+function GNNGraphs._degree((s, t)::Tuple, ::Type{TracedRNumber{T}}, dir::Symbol,
+                           edge_weight::AbstractVector, num_nodes::Int) where {T}
+    return GNNGraphs._degree((s, t), T, dir, edge_weight, num_nodes)
+end
+
+end # module

GNNGraphs/src/GNNGraphs.jl

@@ -84,6 +84,7 @@ export add_nodes,
        perturb_edges,
        remove_nodes,
        ppr_diffusion,
+       pad_graph,
 # from MLUtils.jl
        batch,
        unbatch,

GNNGraphs/src/abstracttypes.jl

@@ -1,10 +1,13 @@
 
-const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V <: Union{Nothing, AbstractVector}}
+# The index vectors of COO_T and AVecI are not restricted to integer eltypes so that
+# traced arrays (e.g. Reactant's `TracedRArray{Int}`, whose eltype is `TracedRNumber{Int}`)
+# can be stored in a graph.
+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}
+const AVecI = AbstractVector
 
 # All concrete graph types should be subtypes of AbstractGNNGraph{T}.
 # GNNGraph and GNNHeteroGraph are the two concrete types.

GNNGraphs/src/convert.jl

@@ -125,15 +125,15 @@ end
 function _to_coo_graph(g::GNNGraph{<:SPARSE_T})
     s, t, v = findnz(g.graph)
     return GNNGraph((s, t, v), g.num_nodes, g.num_edges, g.num_graphs,
-                    g.graph_indicator, g.ndata, g.edata, g.gdata)
+                    g.graph_indicator, g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
 end
 
 function _to_coo_graph(g::GNNGraph{<:ADJMAT_T})
     A = g.graph
     s, t, _ = _findnz_idx(A)
     v = _edge_values(A, s, t)
     return GNNGraph((s, t, v), g.num_nodes, g.num_edges, g.num_graphs,
-                    g.graph_indicator, g.ndata, g.edata, g.gdata)
+                    g.graph_indicator, g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
 end
 
 ### CONVERT TO ADJACENCY MATRIX ################

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

@@ -29,7 +29,8 @@ as well, unless explicitly set by the keyword arguments
 A `GNNGraph` can also represent multiple graphs batched together
 (see [`MLUtils.batch`](@ref) or [`SparseArrays.blockdiag`](@ref)).
 The field `g.graph_indicator` contains the graph membership
-of each node.
+of each node, and `g.graph_sizes` the number of nodes in each graph.
+Both are `nothing` for a single graph.
 
 `GNNGraph`s are always directed graphs, therefore each edge is defined
 by a source node and a target node (see [`edge_index`](@ref)).
@@ -110,23 +111,33 @@ struct GNNGraph{T <: Union{COO_T, ADJMAT_T}} <: AbstractGNNGraph{T}
     num_edges::Int
     num_graphs::Int
     graph_indicator::Union{Nothing, AVecI}       # vector of ints or nothing
+    graph_sizes::Union{Nothing, AVecI}           # number of nodes in each graph, or nothing
     ndata::DataStore
     edata::DataStore
     gdata::DataStore
     is_coalesced::Bool # only for :coo, true if the graph is coalesced, i.e., no multi edges and indices ordered by target, then source
 end
 
-# GNNGraph constructor setting the is_coalesced field to false
+# GNNGraph constructor deriving the graph_sizes field from the graph indicator.
+# graph_sizes duplicates the graph indicator, but its length is the number of graphs,
+# which makes the number of graphs part of the array shapes seen by compilers like Reactant.
 function GNNGraph(graph::T,
                   num_nodes::Int,
                   num_edges::Int,
                   num_graphs::Int,
                   graph_indicator::Union{Nothing, AVecI},
                   ndata::DataStore,
                   edata::DataStore,
-                  gdata::DataStore) where {T <: Union{COO_T, ADJMAT_T}}
+                  gdata::DataStore,
+                  is_coalesced::Bool = false;
+                  graph_sizes = _graph_sizes(graph_indicator, num_graphs)) where {T <: Union{COO_T, ADJMAT_T}}
     return GNNGraph{T}(graph, num_nodes, num_edges, num_graphs,
-                       graph_indicator, ndata, edata, gdata, false)
+                       graph_indicator, graph_sizes, ndata, edata, gdata, is_coalesced)
+end
+
+_graph_sizes(graph_indicator::Nothing, num_graphs) = nothing
+function _graph_sizes(graph_indicator::AVecI, num_graphs)
+    return _scatter(+, ones_like(graph_indicator), graph_indicator, num_graphs)
 end
 
 function GNNGraph(data::D;
@@ -233,7 +244,7 @@ function GNNGraph(g::GNNGraph; ndata = g.ndata, edata = g.edata, gdata = g.gdata
     return GNNGraph(graph,
                 g.num_nodes, g.num_edges, g.num_graphs,
                 g.graph_indicator,
-                ndata, edata, gdata)
+                ndata, edata, gdata; graph_sizes = g.graph_sizes)
 end
 
 """
@@ -247,12 +258,12 @@ function Base.copy(g::GNNGraph; deep = false)
         GNNGraph(deepcopy(g.graph),
                  g.num_nodes, g.num_edges, g.num_graphs,
                  deepcopy(g.graph_indicator),
-                 deepcopy(g.ndata), deepcopy(g.edata), deepcopy(g.gdata))
+                 deepcopy(g.ndata), deepcopy(g.edata), deepcopy(g.gdata); graph_sizes = deepcopy(g.graph_sizes))
     else
         GNNGraph(g.graph,
                  g.num_nodes, g.num_edges, g.num_graphs,
                  g.graph_indicator,
-                 g.ndata, g.edata, g.gdata)
+                 g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
     end
 end
 
@@ -345,14 +356,14 @@ MLUtils.getobs(g::GNNGraph, i) = getgraph(g, i)
 function Base.:(==)(g1::GNNGraph, g2::GNNGraph)
     g1 === g2 && return true
     for k in fieldnames(typeof(g1))
-        k === :graph_indicator && continue
+        k in (:graph_indicator, :graph_sizes) && continue
         getfield(g1, k) != getfield(g2, k) && return false
     end
     return true
 end
 
 function Base.hash(g::T, h::UInt) where {T <: GNNGraph}
-    fs = (getfield(g, k) for k in fieldnames(T) if k !== :graph_indicator)
+    fs = (getfield(g, k) for k in fieldnames(T) if k ∉ (:graph_indicator, :graph_sizes))
     return foldl((h, f) -> hash(f, h), fs, init = hash(T, h))
 end
 

GNNGraphs/src/sampling.jl

@@ -202,7 +202,7 @@ function sample_neighbors(g::GNNGraph{<:COO_T}, nodes, K = -1;
         gnew = GNNGraph(graph,
                         g.num_nodes, num_edges, g.num_graphs,
                         g.graph_indicator,
-                        g.ndata, edata, g.gdata)
+                        g.ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
     else
         nodes_other = dir == :in ? setdiff(s, nodes) : setdiff(t, nodes)
         nodes_all = [nodes; nodes_other]