add GlobalAttentionPool

Author: CarloLucibello

src/GraphNeuralNetworks.jl

@@ -60,6 +60,7 @@ export
 
     # layers/pool
     GlobalPool,
+    GlobalAttentionPool,
     TopKPool,
     topk_index
 

src/layers/pool.jl

@@ -42,6 +42,47 @@ end
 
 (l::GlobalPool)(g::GNNGraph) = GNNGraph(g, gdata=l(g, node_features(g)))
 
+
+@doc raw"""
+    GlobalAttentionPool(fgate, ffeat=identity)
+
+Global soft attention layer from the [Gated Graph Sequence Neural
+Networks](https://arxiv.org/abs/1511.05493) paper
+
+```math
+\mathbf{u}_V} = \sum_{i\in V} \mathrm{softmax} \left(
+                    f_{\mathrm{gate}} ( \mathbf{x}_i ) \right) \odot
+                    f_{\mathrm{feat}} ( \mathbf{x}_i ),
+```
+
+where ``f_{\mathrm{gate}} \colon \mathbb{R}^F \to
+\mathbb{R}`` and ``f_{\mathbf{feat}}` denote neural networks.
+
+# Arguments
+
+fgate: 
+ffeat: 
+"""
+struct GlobalAttentionPool{G,F}
+    fgate::G
+    ffeat::F
+end
+
+@functor GlobalAttentionPool
+
+GlobalAttentionPool(fgate) = GlobalAttentionPool(fgate, identity)
+
+
+function (l::GlobalAttentionPool)(g::GNNGraph, x::AbstractArray)
+    weights = softmax_nodes(g, l.fgate(x))
+    feats = l.ffeat(x)
+    u = reduce_nodes(+, g, weights .* feats)
+    return u   
+end
+
+(l::GlobalAttentionPool)(g::GNNGraph) = GNNGraph(g, gdata=l(g, node_features(g)))
+
+
 """
     TopKPool(adj, k, in_channel)
 
@@ -60,6 +101,9 @@ struct TopKPool{T,S}
     Ã::AbstractMatrix{T}
 end
 
+
+
+
 function TopKPool(adj::AbstractMatrix, k::Int, in_channel::Int; init=glorot_uniform)
     TopKPool(adj, k, init(in_channel), similar(adj, k, k))
 end

test/layers/pool.jl

@@ -1,14 +1,37 @@
 @testset "pool" begin
     @testset "GlobalPool" begin
+        p = GlobalPool(+)
         n = 10
-        X = rand(Float32, 16, n)
+        chin = 6
+        X = rand(Float32, 6, n)
         g = GNNGraph(random_regular_graph(n, 4), ndata=X)
-        p = GlobalPool(+)
-        y = p(g, X)
-        @test y ≈ NNlib.scatter(+, X, ones(Int, n))
+        u = p(g, X)
+        @test u ≈ sum(X, dims=2)
+
+        ng = 3
+        g = Flux.batch([GNNGraph(random_regular_graph(n, 4), 
+                                 ndata=rand(Float32, chin, n)) 
+                        for i=1:ng])
+        u = p(g, g.ndata.x)
+        @test size(u) == (chin, ng)
+        @test u[:,[1]] ≈ sum(g.ndata.x[:,1:n], dims=2)
+        @test p(g).gdata.u == u
+      
         test_layer(p, g, rtol=1e-5, exclude_grad_fields = [:aggr], outtype=:graph)
     end
 
+    @testset "GlobalAttentionPool" begin
+        n = 10
+        chin = 16
+        X = rand(Float32, chin, n)
+        g = GNNGraph(random_regular_graph(n, 4), ndata=X)
+        fgate = Dense(chin, 1, sigmoid)
+        p = GlobalAttentionPool(fgate)
+        y = p(g, X)
+        test_layer(p, g, rtol=1e-5, outtype=:graph)
+    end
+
+
     @testset "TopKPool" begin
         N = 10
         k, in_channel = 4, 7