handle isolated nodes

Author: CarloLucibello

src/GNNGraphs/gatherscatter.jl

@@ -3,11 +3,15 @@ _gather(x::Tuple, i) = map(x -> _gather(x, i), x)
 _gather(x::AbstractArray, i) = NNlib.gather(x, i)
 _gather(x::Nothing, i) = nothing
 
+_scatter(aggr, m::Nothing, t; dstsize=nothing) = nothing
 _scatter(aggr, m::NamedTuple, t; dstsize=nothing) = map(m -> _scatter(aggr, m, t; dstsize), m)
 _scatter(aggr, m::Tuple, t; dstsize=nothing) = map(m -> _scatter(aggr, m, t; dstsize), m)
-_scatter(aggr, m::AbstractArray, t; dstsize=nothing) = NNlib.scatter(aggr, m, t; dstsize)
-_scatter(aggr, m::Nothing, t; dstsize=nothing) = nothing
-
+function _scatter(aggr, m::AbstractMatrix, t; dstsize=nothing)
+    if dstsize !== nothing
+        dstsize = (size(m,1), dstsize)
+    end
+    NNlib.scatter(aggr, m, t; dstsize)
+end
 ## TO MOVE TO NNlib ######################################################
 
 

src/msgpass.jl

@@ -134,7 +134,7 @@ where it comes after [`apply_edges`](@ref).
 """
 function aggregate_neighbors(g::GNNGraph, aggr, m)
     s, t = edge_index(g)
-    return GNNGraphs._scatter(aggr, m, t)
+    return GNNGraphs._scatter(aggr, m, t; dstsize=g.num_nodes)
 end
 
 

test/msgpass.jl

@@ -28,8 +28,14 @@
         m = propagate(message, g, +, xj=X)
 
         @test size(m) == (out_channel, num_V)
-    end
 
+        @testset "isolated nodes" begin
+            x1 = rand(1, 6)
+            g1 = GNNGraph(collect(1:5), collect(1:5), num_nodes=6)
+            y1 = propagate((xi,xj,e) -> xj, g, +, xj=x1)
+            @test size(y1) == (1, 6)
+        end
+    end
 
     @testset "apply_edges" begin
         m = apply_edges(g, e=E) do xi, xj, e
@@ -85,7 +91,7 @@
         @test spmm_copyxj_fused(g) ≈ X * Adj
     end
 
-    @testset "e_mul_xj adn w_mul_xj for weighted conv" begin
+    @testset "e_mul_xj and w_mul_xj for weighted conv" begin
         n = 128
         A = sprand(n, n, 0.1)
         Adj = map(x -> x > 0 ? 1 : 0, A)