fix

Author: CarloLucibello

src/GNNGraphs/gatherscatter.jl

@@ -3,15 +3,20 @@ _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)
-function _scatter(aggr, m::AbstractMatrix, t; dstsize=nothing)
-    if dstsize !== nothing
-        dstsize = (size(m,1), dstsize)
-    end
-    NNlib.scatter(aggr, m, t; dstsize)
+_scatter(aggr, src::Nothing, idx, n) = nothing
+_scatter(aggr, src::NamedTuple, idx, n) = map(s -> _scatter(aggr, s, idx, n), src)
+_scatter(aggr, src::Tuple, idx, n) = map(s -> _scatter(aggr, s, idx, n), src)
+
+function _scatter(aggr, 
+                 src::AbstractArray, 
+                 idx::AbstractVector{<:Integer}, 
+                 n::Integer)
+
+    dstsize =  (size(src)[1:end-1]..., n)
+    NNlib.scatter(aggr, src, idx; dstsize)
 end
+
+
 ## TO MOVE TO NNlib ######################################################
 
 

src/GNNGraphs/transform.jl

@@ -109,8 +109,8 @@ function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr=+)
         idxs .= 1:num_edges
         idxs .= idxs .- cumsum(.!mask)
         num_edges = length(s)
-        w = _scatter(aggr, w, idxs)
-        edata = _scatter(aggr, edata, idxs)
+        w = _scatter(aggr, w, idxs, num_edges)
+        edata = _scatter(aggr, edata, idxs, num_edges)
     end
 
     GNNGraph((s, t, w),

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; dstsize=g.num_nodes)
+    return GNNGraphs._scatter(aggr, m, t, g.num_nodes)
 end