Add GlobalAttentionPool

Author: aurorarossi

GNNLux/src/GNNLux.jl

@@ -49,7 +49,8 @@ export TGCN,
        EvolveGCNO
 
 include("layers/pool.jl")
-export GlobalPool
+export GlobalPool,
+       GlobalAttentionPool
 
 end #module
 

GNNLux/src/layers/pool.jl

@@ -39,4 +39,71 @@ end
 
 (l::GlobalPool)(g::GNNGraph, x::AbstractArray, ps, st) = GNNlib.global_pool(l, g, x), st
 
-(l::GlobalPool)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g), ps, st))
+(l::GlobalPool)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g), ps, st))
+
+@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} \alpha_i\, f_{feat}(\mathbf{x}_i)
+```
+
+where the coefficients ``\alpha_i`` are given by a [`softmax_nodes`](@ref)
+operation:
+
+```math
+\alpha_i = \frac{e^{f_{gate}(\mathbf{x}_i)}}
+                {\sum_{i'\in V} e^{f_{gate}(\mathbf{x}_{i'})}}.
+```
+
+# Arguments
+
+- `fgate`: The function ``f_{gate}: \mathbb{R}^{D_{in}} \to \mathbb{R}``. 
+           It is typically expressed by a neural network.
+
+- `ffeat`: The function ``f_{feat}: \mathbb{R}^{D_{in}} \to \mathbb{R}^{D_{out}}``. 
+           It is typically expressed by a neural network.
+
+# Examples
+
+```julia
+using Graphs, LuxCore, Lux, GNNLux, Random
+
+rng = Random.default_rng()
+chin = 6
+chout = 5    
+
+fgate = Dense(chin, 1)
+ffeat = Dense(chin, chout)
+pool = GlobalAttentionPool(fgate, ffeat)
+
+g = batch([GNNGraph(Graphs.random_regular_graph(10, 4), 
+                         ndata=rand(Float32, chin, 10)) 
+                for i=1:3])
+
+ps = (fgate = LuxCore.initialparameters(rng, fgate), ffeat = LuxCore.initialparameters(rng, ffeat))
+st = (fgate = LuxCore.initialstates(rng, fgate), ffeat = LuxCore.initialstates(rng, ffeat))
+
+u, st = pool(g, g.ndata.x, ps, st)
+
+@assert size(u) == (chout, g.num_graphs)
+```
+"""
+struct GlobalAttentionPool{G, F}
+    fgate::G
+    ffeat::F
+end
+
+GlobalAttentionPool(fgate) = GlobalAttentionPool(fgate, identity)
+
+function (l::GlobalAttentionPool)(g, x, ps, st)
+    fgate = StatefulLuxLayer{true}(l.fgate, ps.fgate, _getstate(st, :fgate))
+    ffeat = StatefulLuxLayer{true}(l.ffeat, ps.ffeat, _getstate(st, :ffeat))
+    m = (; fgate, ffeat)
+    return GNNlib.global_attention_pool(m, g, x), st
+end
+
+(l::GlobalAttentionPool)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g), ps, st))

GNNLux/test/layers/pool.jl

@@ -1,15 +1,21 @@
 @testitem "Pooling" setup=[TestModuleLux] begin
     using .TestModuleLux
-    @testset "GlobalPool" begin
+    @testset "Pooling" begin
 
         rng = StableRNG(1234)
         g = rand_graph(rng, 10, 40)
         in_dims = 3
         x = randn(rng, Float32, in_dims, 10)
 
-        @testset "GCNConv" begin
-            l = GlobalPool(mean)
+        @testset "GlobalPool" begin
+            l = GNNLux.GlobalPool(mean)
             test_lux_layer(rng, l, g, x, sizey=(in_dims,1))
         end
+        @testset "GlobalAttentionPool" begin
+            fgate = Dense(in_dims, 1)
+            ffeat = Dense(in_dims, in_dims)
+            l = GNNLux.GlobalAttentionPool(fgate, ffeat)
+            test_lux_layer(rng, l, g, x, sizey=(in_dims,1), container=true)
+        end
     end
 end

GraphNeuralNetworks/src/layers/pool.jl

@@ -61,10 +61,10 @@ operation:
 # Arguments
 
 - `fgate`: The function ``f_{gate}: \mathbb{R}^{D_{in}} \to \mathbb{R}``. 
-           It is tipically expressed by a neural network.
+           It is typically expressed by a neural network.
 
 - `ffeat`: The function ``f_{feat}: \mathbb{R}^{D_{in}} \to \mathbb{R}^{D_{out}}``. 
-           It is tipically expressed by a neural network.
+           It is typically expressed by a neural network.
 
 # Examples