better handling of chain with graph input

Author: CarloLucibello

src/layers/basic.jl

@@ -60,6 +60,11 @@ end
 
 Flux.functor(::Type{<:GNNChain}, c) = c.layers, ls -> GNNChain(ls...)
 
+# input from graph
+applylayer(l, g::GNNGraph) = GNNGraph(g, ndata=l(node_features(g)))
+applylayer(l::GNNLayer, g::GNNGraph) = l(g)
+
+# explicit input
 applylayer(l, g::GNNGraph, x) = l(x)
 applylayer(l::GNNLayer, g::GNNGraph, x) = l(g, x)
 
@@ -68,11 +73,17 @@ applylayer(l::Parallel, g::GNNGraph, x::AbstractArray) = mapreduce(f -> applylay
 applylayer(l::Parallel, g::GNNGraph, xs::Vararg{<:AbstractArray}) = mapreduce((f, x) -> applylayer(f, g, x), l.connection, l.layers, xs)
 applylayer(l::Parallel, g::GNNGraph, xs::Tuple) = applylayer(l, g, xs...)
 
+# input from graph
+applychain(::Tuple{}, g::GNNGraph) = g
+applychain(fs::Tuple, g::GNNGraph) = applychain(tail(fs), applylayer(first(fs), g))
 
+# explicit input
 applychain(::Tuple{}, g::GNNGraph, x) = x
 applychain(fs::Tuple, g::GNNGraph, x) = applychain(tail(fs), g, applylayer(first(fs), g, x))
 
 (c::GNNChain)(g::GNNGraph, x) = applychain(Tuple(c.layers), g, x)
+(c::GNNChain)(g::GNNGraph) = applychain(Tuple(c.layers), g)
+
 
 Base.getindex(c::GNNChain, i::AbstractArray) = GNNChain(c.layers[i]...)
 Base.getindex(c::GNNChain{<:NamedTuple}, i::AbstractArray) = 

test/layers/basic.jl

@@ -31,6 +31,18 @@
 
             test_layer(gnn, g, rtol=1e-5, exclude_grad_fields=[:μ, :σ²])
         end
+
+        @testset "Only graph input" begin
+            nin, nout = 2, 4
+            ndata = rand(nin, 3)
+            edata = rand(nin, 3)
+            g = GNNGraph([1,1,2], [2, 3, 3], ndata=ndata, edata=edata)
+            m = NNConv(nin => nout, Dense(2, nin*nout, tanh))
+            chain = GNNChain(m)
+            y = m(g, g.ndata.x, g.edata.e)
+            @test m(g).ndata.x == y
+            @test chain(g).ndata.x == y
+        end
     end
 end