add Set2Set (#276)

* add Set2Set

* cleanup

* cleanup

Author: CarloLucibello

src/GNNGraphs/transform.jl

@@ -729,3 +729,4 @@ ci2t(ci::AbstractVector{<:CartesianIndex}, dims) = ntuple(i -> map(x -> x[i], ci
 @non_differentiable negative_sample(x...)
 @non_differentiable add_self_loops(x...)     # TODO this is wrong, since g carries feature arrays, needs rrule
 @non_differentiable remove_self_loops(x...)  # TODO this is wrong, since g carries feature arrays, needs rrule
+@non_differentiable dense_zeros_like(x...)

src/GraphNeuralNetworks.jl

@@ -7,6 +7,7 @@ using CUDA
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor, batch
 using MacroTools: @forward
+using MLUtils
 using NNlib, NNlibCUDA
 using NNlib: scatter, gather
 using ChainRulesCore
@@ -69,6 +70,7 @@ export
 # layers/pool
       GlobalPool,
       GlobalAttentionPool,
+      Set2Set,
       TopKPool,
       topk_index,
 

src/layers/pool.jl

@@ -143,3 +143,62 @@ function topk_index(y::AbstractVector, k::Int)
 end
 
 topk_index(y::Adjoint, k::Int) = topk_index(y', k)
+
+
+@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} <: GNNLayer
+    lstm::L
+    num_iters::Int
+end
+
+@functor Set2Set
+
+function Set2Set(n_in::Int, n_iters::Int, n_layers::Int = 1)
+    @assert n_layers >= 1
+    n_out = 2 * n_in
+
+    if n_layers == 1
+        lstm = LSTM(n_out => n_in)
+    else
+        layers = [LSTM(n_out => n_in)]
+        for _ in 2:n_layers
+            push!(layers, LSTM(n_in => n_in))
+        end
+        lstm = Chain(layers...)
+    end
+
+    return Set2Set(lstm, n_iters)
+end
+
+function (l::Set2Set)(g::GNNGraph, x::AbstractMatrix)
+    n_in = size(x, 1)
+    Flux.reset!(l.lstm)
+    qstar = zeros_like(x, (2*n_in, g.num_graphs))
+    for t in 1:l.num_iters
+        q = l.lstm(qstar)                            # [n_in, n_graphs]
+        qn = broadcast_nodes(g, q)                    # [n_in, n_nodes]
+        α = softmax_nodes(g, sum(qn .* x, dims = 1))  # [1, n_nodes]
+        r = reduce_nodes(+, g, x .* α)               # [n_in, n_graphs]
+        qstar = vcat(q, r)                           # [2*n_in, n_graphs]
+    end
+    return qstar
+end
+
+(l::Set2Set)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g)))

test/layers/pool.jl

@@ -60,3 +60,18 @@ end
     @test topk_index(X, 4) == [1, 2, 3, 4]
     @test topk_index(X', 4) == [1, 2, 3, 4]
 end
+
+@testset "Set2Set" begin
+    n_in = 3
+    n_iters = 2
+    n_layers = 1
+    g = batch([rand_graph(10, 40, graph_type = GRAPH_T) for _ in 1:5])
+    g = GNNGraph(g, ndata = rand(Float32, n_in, g.num_nodes))
+    l = Set2Set(n_in, n_iters, n_layers)
+    y = l(g, node_features(g))
+    @test size(y) == (2 * n_in, g.num_graphs)
+    
+    ## TODO the numerical gradient seems to be 3 times smaller than zygote one
+    # test_layer(l, g, rtol = 1e-4, atol=1e-4, outtype = :graph, outsize = (2 * n_in, g.num_graphs), 
+    #         verbose=true, exclude_grad_fields = [:state0, :state])
+end

test/test_utils.jl

@@ -180,8 +180,10 @@ function test_approx_structs(l, l̄, l̄fd; atol = 1e-5, rtol = 1e-5,
 
     for f in fieldnames(typeof(l))
         f ∈ exclude_grad_fields && continue
+        verbose && println("Test gradient of field $f...")
         x, g, gfd = getfield(l, f), getfield(l̄, f), getfield(l̄fd, f)
         test_approx_structs(x, g, gfd; atol, rtol, exclude_grad_fields, verbose)
+        verbose && println("... field $f done!")
     end
     return true
 end