better printing and chech feature types

Author: CarloLucibello

prova.jl

@@ -0,0 +1,20 @@
+using GraphNeuralNetworks, Test, Flux
+using BenchmarkTools
+using ProfileView
+
+g = rand_graph(10, 30, ndata=rand(Float32, 2, 10))
+l = GATConv(2 => 2)
+y = l(g, g.ndata.x)
+@assert eltype(y) == Float32
+
+dx = gradient(x -> sum(sin.(l(g, x))), g.ndata.x)[1]
+@assert eltype(dx) == Float32
+
+struct B
+    slope
+end
+
+(a::B)(x) = leakyrelu(x, a.slope)
+
+a = B(0.3f0)
+grad = gradient(a -> a(0.1f0), a)[1]

prova2.jl

@@ -0,0 +1,26 @@
+using Flux
+using GraphNeuralNetworks
+using BenchmarkTools
+using ProfileView
+
+f(x) = 1
+
+function test(g)
+    loader = Flux.DataLoader(g, batchsize=100, shuffle=true)
+    s = 0 
+    for d in loader
+        s += f(d)
+    end
+    return s
+end
+
+n = 5000
+s = 10
+data = [rand_graph(s, s, ndata = rand(1, s)) for i in 1:n] 
+x1 = Flux.batch(data) 
+x2 = Flux.batch([rand(s + s + s + s) for i in 1:n]) #source+target+data+extra
+
+# @profview test(x1)
+@btime test($x1);   #  1.295 s (2502 allocations: 6.17 MiB)
+@btime test($x2);   #  357.595 μs (152 allocations: 1.61 MiB)
+@btime test($data); #  65.288 ms (227002 allocations: 27.00 MiB) # this PR

src/GNNGraphs/gnngraph.jl

@@ -206,19 +206,20 @@ function Base.show(io::IO, g::GNNGraph)
     if !isempty(g.ndata)
         print(io, "\n    ndata:")
         for k in keys(g.ndata)
-            print(io, "\n        $k => $(size(g.ndata[k]))")
+            # print(io, "\n        $k => $(size(g.ndata[k]))")
+            print(io, "\n        $k => $(summary(g.ndata[k]))")
         end
     end
     if !isempty(g.edata)
         print(io, "\n    edata:")
         for k in keys(g.edata)
-            print(io, "\n        $k => $(size(g.edata[k]))")
+            print(io, "\n        $k => $(summary(g.edata[k]))")
         end
     end
     if !isempty(g.gdata)
         print(io, "\n    gdata:")
         for k in keys(g.gdata)
-            print(io, "\n        $k => $(size(g.gdata[k]))")
+            print(io, "\n        $k => $(summary(g.gdata[k]))")
         end
     end
 end

src/GNNGraphs/utils.jl

@@ -57,10 +57,14 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
     # This had to workaround two Zygote bugs with NamedTuples
     # https://github.com/FluxML/Zygote.jl/issues/1071
     # https://github.com/FluxML/Zygote.jl/issues/1072
+
+    if n != 1
+        @assert all(x -> x isa AbstractArray, data) "Non-array features provided."
+    end
 
     if n == 1
         # If last array dimension is not 1, add a new dimension. 
-        # This is mostly usefule to reshape globale feature vectors
+        # This is mostly useful to reshape global feature vectors
         # of size D to Dx1 matrices.
         function unsqz(v)
             if v isa AbstractArray && size(v)[end] != 1
@@ -73,19 +77,20 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
     end
 
     sz = map(x -> x isa AbstractArray ? size(x)[end] : 0, data)
+
     if duplicate_if_needed 
         # Used to copy edge features on reverse edges    
         @assert all(s -> s == 0 ||  s == n || s == n÷2, sz)  "Wrong size in last dimension for feature array."
 
         function duplicate(v)
-            if v isa AbstractArray && size(v)[end] == n÷2
+            if size(v)[end] == n÷2
                 v = cat(v, v, dims=ndims(v))
             end
             v
         end
         data = NamedTuple{keys(data)}(duplicate.(values(data)))
     else
-        @assert all(s -> s == 0 ||  s == n, sz) "Wrong size in last dimension for feature array."
+        @assert all(x -> x == 0 || x == n, sz) "Wrong size in last dimension for feature array."
     end
     return data
 end

test/GNNGraphs/gnngraph.jl

@@ -229,6 +229,11 @@
         # Attach non array data
         g = GNNGraph(erdos_renyi(10,  30), edata="ciao", graph_type=GRAPH_T)
         @test g.edata.e == "ciao"
+
+
+        # Wrong need number of features
+        g = GNNGraph(erdos_renyi(10,  30), edata="ciao", graph_type=GRAPH_T)
+
     end
 
     @testset "LearnBase and DataLoader compat" begin