Fixes

Author: aurorarossi

GNNLux/src/layers/pool.jl

@@ -51,7 +51,7 @@ Networks](https://arxiv.org/abs/1511.05493) paper
 \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)
+where the coefficients ``\alpha_i`` are given by a [`GNNLib.softmax_nodes`](@ref)
 operation:
 
 ```math
@@ -92,9 +92,9 @@ u, st = pool(g, g.ndata.x, ps, st)
 @assert size(u) == (chout, g.num_graphs)
 ```
 """
-struct GlobalAttentionPool{G, F}
-    fgate::G
-    ffeat::F
+struct GlobalAttentionPool <: GNNContainerLayer{(:fgate, :ffeat)}
+    fgate
+    ffeat
 end
 
 GlobalAttentionPool(fgate) = GlobalAttentionPool(fgate, identity)

GNNLux/test/layers/pool.jl

@@ -3,19 +3,26 @@
     @testset "Pooling" begin
 
         rng = StableRNG(1234)
-        g = rand_graph(rng, 10, 40)
-        in_dims = 3
-        x = randn(rng, Float32, in_dims, 10)
-
         @testset "GlobalPool" begin
+            g = rand_graph(rng, 10, 40)
+            in_dims = 3
+            x = randn(rng, Float32, in_dims, 10)
             l = 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)
+            n = 10
+            chin = 6
+            chout = 5
+            ng = 3
+            g = batch([GNNGraph(rand_graph(rng, 10, 40),
+                            ndata = rand(Float32, chin, n)) for i in 1:ng])
+
+            fgate = Dense(chin, 1)
+            ffeat = Dense(chin, chout)
             l = GlobalAttentionPool(fgate, ffeat)
-            test_lux_layer(rng, l, g, x, sizey=(in_dims,1), container=true)
+    
+            test_lux_layer(rng, l, g, g.ndata.x, sizey=(chout,ng), container=true)
         end
     end
 end