Add TopK pooling

Author: aurorarossi

GNNLux/src/layers/pool.jl

@@ -106,4 +106,66 @@ function (l::GlobalAttentionPool)(g, x, ps, st)
     return GNNlib.global_attention_pool(m, g, x), st
 end
 
-(l::GlobalAttentionPool)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g), ps, st))
+(l::GlobalAttentionPool)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g), ps, st))
+
+"""
+    TopKPool(adj, k, in_channel)
+
+Top-k pooling layer.
+
+# Arguments
+
+- `adj`: Adjacency matrix of a graph.
+- `k`: Top-k nodes are selected to pool together.
+- `in_channel`: The dimension of input channel.
+"""
+struct TopKPool{T, S}
+    A::AbstractMatrix{T}
+    k::Int
+    p::AbstractVector{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
+
+(t::TopKPool)(x::AbstractArray, ps, st) = GNNlib.topk_pool(t, x)
+
+
+@doc raw"""
+    Set2Set(n_in, n_iters, n_layers = 1)
+
+Set2Set layer from the paper [Order Matters: Sequence to sequence for sets](https://arxiv.org/abs/1511.06391).
+
+For each graph in the batch, the layer computes an output vector of size `2*n_in` by iterating the following steps `n_iters` times:
+```math
+\mathbf{q} = \mathrm{LSTM}(\mathbf{q}_{t-1}^*)
+\alpha_{i} = \frac{\exp(\mathbf{q}^T \mathbf{x}_i)}{\sum_{j=1}^N \exp(\mathbf{q}^T \mathbf{x}_j)} 
+\mathbf{r} = \sum_{i=1}^N \alpha_{i} \mathbf{x}_i
+\mathbf{q}^*_t = [\mathbf{q}; \mathbf{r}]
+```
+where `N` is the number of nodes in the graph, `LSTM` is a Long-Short-Term-Memory network with `n_layers` layers, input size `2*n_in` and output size `n_in`.
+
+Given a batch of graphs `g` and node features `x`, the layer returns a matrix of size `(2*n_in, n_graphs)`.
+```
+"""
+struct Set2Set{L} <: GNNContainerLayer{(:lstm,)}
+    lstm::L
+    num_iters::Int
+end
+
+function Set2Set(n_in::Int, n_iters::Int, n_layers::Int = 1)
+    @assert n_layers == 1 "multiple layers not implemented yet" #TODO
+    n_out = 2 * n_in
+    lstm = Lux.LSTMCell(n_out => n_in)
+    return Set2Set(lstm, n_iters)
+end
+
+function (l::Set2Set)(g, x, ps, st)
+    lstm = StatefulLuxLayer{true}(l.lstm, ps.lstm, _getstate(st, :lstm))
+    m = (; lstm, Wh = ps.lstm.weight_hh)
+    return GNNlib.set2set_pool(m, g, x)
+end
+
+(l::Set2Set)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g), ps, st))

GNNLux/test/layers/pool.jl

@@ -24,5 +24,23 @@
 
             test_lux_layer(rng, l, g, g.ndata.x, sizey=(chout,ng), container=true)
         end
+
+        @testset "TopKPool" begin
+            N = 10
+            k, in_channel = 4, 7
+            X = rand(in_channel, N)
+            ps = (;)
+            st = (;)
+            for T in [Bool, Float64]
+                adj = rand(T, N, N)
+                p = GNNLux.TopKPool(adj, k, in_channel)
+                @test eltype(p.p) === Float32
+                @test size(p.p) == (in_channel,)
+                @test eltype(p.Ã) === T
+                @test size(p.Ã) == (k, k)
+                y = p(X, ps, st)
+                @test size(y) == (in_channel, k)
+            end
+        end
     end
 end