generalize propagation to tuple and namedtuple

Author: CarloLucibello

src/gnngraph.jl

@@ -59,7 +59,7 @@ from the LightGraphs' graph library can be used on it.
 
 # Usage. 
 
-```
+```julia
 using Flux, GraphNeuralNetworks
 
 # Construct from adjacency list representation

src/layers/conv.jl

@@ -196,7 +196,7 @@ end
 
 
 @doc raw"""
-    GATConv(in => out;
+    GATConv(in => out, , σ=identity;
             heads=1,
             concat=true,
             init=glorot_uniform    
@@ -228,6 +228,7 @@ struct GATConv{T, A<:AbstractMatrix{T}, B} <: GNNLayer
     weight::A
     bias::B
     a::A
+    σ
     negative_slope::T
     channel::Pair{Int, Int}
     heads::Int
@@ -237,44 +238,43 @@ end
 @functor GATConv
 Flux.trainable(l::GATConv) = (l.weight, l.bias, l.a)
 
-function GATConv(ch::Pair{Int,Int};
+function GATConv(ch::Pair{Int,Int}, σ=identity;
                  heads::Int=1, concat::Bool=true, negative_slope=0.2f0,
                  init=glorot_uniform, bias::Bool=true)
     in, out = ch             
     W = init(out*heads, in)
     b = Flux.create_bias(W, bias, out*heads)
     a = init(2*out, heads)
-    GATConv(W, b, a, negative_slope, ch, heads, concat)
+    GATConv(W, b, a, σ, negative_slope, ch, heads, concat)
 end
 
-function (gat::GATConv)(g::GNNGraph, X::AbstractMatrix)
-    check_num_nodes(g, X)
+function compute_message(l::GATConv, Wxi, Wxj)
+    aWW = sum(l.a .* cat(Wxi, Wxj, dims=1), dims=1)   # 1 × nheads × nedges
+    α = exp.(leakyrelu.(aWW, l.negative_slope))       
+    return (α = α, m = α .* Wxj)
+end
+
+update_node(l::GATConv, d̄, x) = d̄.m ./ d̄.α
+
+function (l::GATConv)(g::GNNGraph, x::AbstractMatrix)
+    check_num_nodes(g, x)
     g = add_self_loops(g)
-    chin, chout = gat.channel
-    heads = gat.heads
+    chin, chout = l.channel
+    heads = l.heads
 
-    source, target = edge_index(g)
-    Wx = gat.weight*X
+    Wx = l.weight * x
     Wx = reshape(Wx, chout, heads, :)                   # chout × nheads × nnodes
-    Wxi = NNlib.gather(Wx, target)                      # chout × nheads × nedges
-    Wxj = NNlib.gather(Wx, source)
-
-    # Edge Message
-    # Computing softmax. TODO make it numerically stable
-    aWW = sum(gat.a .* cat(Wxi, Wxj, dims=1), dims=1)   # 1 × nheads × nedges
-    α = exp.(leakyrelu.(aWW, gat.negative_slope))       
-    m̄ =  NNlib.scatter(+, α .* Wxj, target)             # chout × nheads × nnodes 
-    ᾱ = NNlib.scatter(+, α, target)                     # 1 × nheads × nnodes
 
-    # Node update
-    b = reshape(gat.bias, chout, heads)
-    X = m̄ ./ ᾱ .+ b                                     # chout × nheads × nnodes 
-    if !gat.concat
-        X = sum(X, dims=2)
+    x, _ = propagate(l, g, +, Wx)                 ## chout × nheads × nnodes
+
+    b = reshape(l.bias, chout, heads)
+    x = l.σ.(x .+ b)                                      
+    if !l.concat
+        x = sum(x, dims=2)
     end
 
     # We finally return a matrix
-    return reshape(X, :, size(X, 3)) 
+    return reshape(x, :, size(x, 3)) 
 end
 
 

src/msgpass.jl

@@ -24,7 +24,7 @@ this method in the forward pass:
 
 # Usage example
 
-```
+```julia
 using GraphNeuralNetworks, Flux
 
 struct GNNConv <: GNNLayer
@@ -101,7 +101,9 @@ function compute_message end
 @inline compute_message(l, x_i, x_j, e_ij) = compute_message(l, x_i, x_j)
 @inline compute_message(l, x_i, x_j) = x_j
 
-_gather(x, i) = NNlib.gather(x, i)
+_gather(x::NamedTuple, i) = map(x -> _gather(x, i), x)
+_gather(x::Tuple, i) = map(x -> _gather(x, i), x)
+_gather(x::AbstractArray, i) = NNlib.gather(x, i)
 _gather(x::Nothing, i) = nothing
 
 function compute_batch_message(l, g, x, e)
@@ -114,9 +116,14 @@ end
 
 ##  Step 2
 
+_scatter(aggr, e::NamedTuple, t) = map(e -> _scatter(aggr, e, t), e)
+_scatter(aggr, e::Tuple, t) = map(e -> _scatter(aggr, e, t), e)
+_scatter(aggr, e::AbstractArray, t) = NNlib.scatter(aggr, e, t)
+_scatter(aggr, e::Nothing, t) = nothing
+
 function aggregate_neighbors(l, g, aggr, e)
     s, t = edge_index(g)
-    NNlib.scatter(aggr, e, t)
+    _scatter(aggr, e, t)
 end
 
 aggregate_neighbors(l, g, aggr::Nothing, e) = nothing

test/msgpass.jl

@@ -1,3 +1,5 @@
+import GraphNeuralNetworks: compute_message, update_node, update_edge, propagate
+
 @testset "message passing" begin 
     in_channel = 10
     out_channel = 5
@@ -113,9 +115,9 @@
         GraphNeuralNetworks.compute_message(l::NewLayerW{GRAPH_T}, x_i, x_j, e_ij) = l.weight * x_j
         GraphNeuralNetworks.update_node(l::NewLayerW{GRAPH_T}, m, x) = l.weight * x + m
 
-        l = NewLayerW(in_channel, out_channel)
         (l::NewLayerW{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)
 
+        l = NewLayerW(in_channel, out_channel)
         g = GNNGraph(adj, ndata=X, edata=E, gdata=U, graph_type=GRAPH_T)
         g_ = l(g)
 
@@ -124,4 +126,37 @@
         @test edge_features(g_) === E
         @test graph_features(g_) === U
     end
+
+    @testset "NamedTuples" begin
+        struct NewLayerNT{G}
+            W
+        end
+        
+        NewLayerNT(in, out) = NewLayerW{GRAPH_T}(randn(T, out, in))
+        
+        function compute_message(l::NewLayerW{GRAPH_T}, di, dj, dij)
+            a = l.W * (di.x .+ dj.x) + dij.e
+            b = l.W * di.x
+            return (; a, b)
+        end
+        function update_node(l::NewLayerW{GRAPH_T}, m, x) 
+            return (α=l.W * x + m.a + m.b, β=m)
+        end
+        function update_edge(l::NewLayerW{GRAPH_T}, m, e) 
+            return m.a
+        end
+
+        function (::NewLayerNT)(l, g, x, e)
+            x, e = propagate(l, g, (; x), (; e))
+            return x.α .+ x.β.a, e
+        end
+
+
+        l = NewLayerNT(in_channel, out_channel)
+        g = GNNGraph(adj, graph_type=GRAPH_T)
+        X′, E′ = l(g, X, E)
+
+        @test size(X′) == (out_channel, num_V)
+        @test size(E′) == (out_channel, num_E)
+    end
 end