more size checks

Author: CarloLucibello

src/GNNGraphs/gatherscatter.jl

@@ -13,7 +13,7 @@ function _scatter(aggr,
                  n::Integer)
 
     dstsize =  (size(src)[1:end-1]..., n)
-    NNlib.scatter(aggr, src, idx; dstsize)
+    return NNlib.scatter(aggr, src, idx; dstsize)
 end
 
 

src/GNNGraphs/utils.jl

@@ -1,9 +1,25 @@
 function check_num_nodes(g::GNNGraph, x::AbstractArray)
-    @assert g.num_nodes == size(x, ndims(x))    
+    @assert g.num_nodes == size(x, ndims(x)) "Got $(size(x, ndims(x))) as last dimension size instead of num_edges=$(g.num_nodes)"
+    return true
 end
+function check_num_nodes(g::GNNGraph, x::NamedTuple)
+    map(x -> check_num_nodes(g, x), x)
+    return true
+end
+
+check_num_nodes(::GNNGraph, ::Nothing) = true
+
 function check_num_edges(g::GNNGraph, e::AbstractArray)
-    @assert g.num_edges == size(e, ndims(e))    
+    @assert g.num_edges == size(e, ndims(e)) "Got $(size(e, ndims(e))) as last dimension size instead of num_edges=$(g.num_edges)"
+    return true    
 end
+function check_num_edges(g::GNNGraph, x::NamedTuple)
+    map(x -> check_num_edges(g, x), x)
+    return true
+end
+
+check_num_edges(::GNNGraph, ::Nothing) = true
+
 
 sort_edge_index(eindex::Tuple) = sort_edge_index(eindex...)
 

src/msgpass.jl

@@ -111,6 +111,9 @@ apply_edges(l, g::GNNGraph; xi=nothing, xj=nothing, e=nothing) =
     apply_edges(l, g, xi, xj, e)
 
 function apply_edges(f, g::GNNGraph, xi, xj, e)
+    check_num_nodes(g, xi)
+    check_num_nodes(g, xj)
+    check_num_edges(g, e)
     s, t = edge_index(g)
     xi = GNNGraphs._gather(xi, t)   # size: (D, num_nodes) -> (D, num_edges)
     xj = GNNGraphs._gather(xj, s)
@@ -133,6 +136,7 @@ Neighborhood aggregation is the second step of [`propagate`](@ref),
 where it comes after [`apply_edges`](@ref).
 """
 function aggregate_neighbors(g::GNNGraph, aggr, m)
+    check_num_edges(g, m)
     s, t = edge_index(g)
     return GNNGraphs._scatter(aggr, m, t, g.num_nodes)
 end

test/msgpass.jl

@@ -62,6 +62,19 @@
         end 
 
         @test m.a == ones(out_channel, num_E)
+
+        @testset "sizecheck" begin
+            x = rand(3, g.num_nodes-1)
+            @test_throws AssertionError apply_edges(copy_xj, g, xj=x)
+            @test_throws AssertionError apply_edges(copy_xj, g, xi=x)
+            
+            x = (a=rand(3, g.num_nodes), b=rand(3, g.num_nodes+1))
+            @test_throws AssertionError apply_edges(copy_xj, g, xj=x)
+            @test_throws AssertionError apply_edges(copy_xj, g, xi=x)
+
+            e = rand(3, g.num_edges-1)
+            @test_throws AssertionError apply_edges(copy_xj, g, e=e)
+        end
     end
 
     @testset "copy_xj" begin