Add WeightAndSum pooling layer

src/GraphNeuralNetworks.jl

@@ -71,6 +71,7 @@ export
       GlobalAttentionPool,
       TopKPool,
       topk_index,
+      WeigthAndSumPool,
 
 # mldatasets
       mldataset2gnngraph

src/layers/pool.jl

@@ -143,3 +143,35 @@ function topk_index(y::AbstractVector, k::Int)
 end
 
 topk_index(y::Adjoint, k::Int) = topk_index(y', k)
+
+"""
+    WeigthAndSumPool(in_feats)
+
+WeigthAndSum sum pooling layer.
+Takes a graph and the node features as inputs, computes the weights for each node and perform a weighted sum.
+"""
+struct WeigthAndSumPool
+    in_feats::Int
+    dense_layer::Dense
+end
+
+@functor WeigthAndSumPool
+
+Flux.trainable(ws::WeigthAndSumPool) = (ws.dense_layer)
+
+function WeigthAndSumPool(in_feats::Int)
+    dense_layer = Dense(in_feats, 1, sigmoid; bias = true)
+    WeigthAndSumPool(in_feats, dense_layer)
+end
+
+function (ws::WeigthAndSumPool)(g::GNNGraph, x::AbstractArray)
+    atom_weighting = ws.dense_layer
+    return reduce_nodes(+, g, atom_weighting(x) .* x)
+end
+
+function (ws::WeigthAndSumPool)(g::GNNGraph, x::CuArray)
+    atom_weighting = ws.dense_layer |> gpu
+    return reduce_nodes(+, g, atom_weighting(x) .* x)
+end
+
+(ws::WeigthAndSumPool)(g::GNNGraph) = GNNGraph(g, gdata = ws(g, node_features(g)))

test/layers/pool.jl

@@ -53,6 +53,22 @@ end
         y = p(X)
         @test size(y) == (in_channel, k)
     end
+
+    @testset "WeigthAndSumPool" begin
+        n = 3
+        chin = 5
+        ng = 3
+
+        ws = WeigthAndSumPool(chin)
+        g = GNNGraph(rand_graph(n, 4), ndata = rand(Float32, chin, n), graph_type = GRAPH_T)
+
+        test_layer(ws, g, rtol = 1e-5, outtype = :graph,outsize = (chin, 1))
+        g_batch = Flux.batch([GNNGraph(rand_graph(n, 4),
+                                       ndata = rand(Float32, chin, n),
+                                       graph_type = GRAPH_T)
+                              for i in 1:ng])
+        test_layer(ws, g_batch, rtol = 1e-5,outtype = :graph, outsize = (chin, ng))
+    end
 end
 
 @testset "topk_index" begin