fix tests

Author: CarloLucibello

docs/src/messagepassing.md

@@ -1,4 +1,21 @@
 # Message Passing
 
-TODO
+The message passing is initiated by [`propagate`](@ref)
+and can be customized for a specific layer by overloading the methods
+[`compute_message`](@ref), [`update_node`](@ref), and [`update_edge`](@ref).
 
+
+The message passing corresponds to the following operations 
+
+```math
+\begin{aligned}
+\mathbf{m}_{j\to i} &= \phi(\mathbf{x}_i, \mathbf{x}_j, \mathbf{e}_{j\to i}) \\
+\mathbf{x}_{i}' &= \gamma_x(\mathbf{x}_{i}, \square_{j\in N(i)}  \mathbf{m}_{j\to i})\\
+\mathbf{e}_{j\to i}^\prime &=  \gamma_e(\mathbf{e}_{j \to i},\mathbf{m}_{j \to i})
+\end{aligned}
+```
+where ``phi`` is expressed by the [`compute_message`](@ref) function, 
+``\gamma_x`` and ``gamma_v`` by [`update_node`](@ref) and [`update_edge`](@ref)
+respectively.
+
+See [`GraphConv`](ref) and [`GATConv`](ref)'s implementations as usage examples. 

src/msgpass.jl

@@ -13,7 +13,7 @@ The computational steps are the following:
 
 ```julia
 m = compute_batch_message(l, g, x, e)  # calls `compute_message`
-m̄ = aggregate_neighbors(l, aggr, g, m)
+m̄ = aggregate_neighbors(g, aggr, m)
 x′ = update_node(l, m̄, x)
 e′ = update_edge(l, m, e)
 ```
@@ -63,7 +63,7 @@ end
 
 function propagate(l, g::GNNGraph, aggr, x, e=nothing)
     m = compute_batch_message(l, g, x, e) 
-    m̄ = aggregate_neighbors(l, g, aggr, m)
+    m̄ = aggregate_neighbors(g, aggr, m)
     x′ = update_node(l, m̄, x)
     e′ = update_edge(l, m, e)
     return x′, e′
@@ -74,15 +74,17 @@ end
 """
     compute_message(l, x_i, x_j, [e_ij])
 
-Message function for the message-passing scheme,
-returning the message from node `j` to node `i` .
+Message function for the message-passing scheme
+started by [`propagate`](@ref).
+Returns the message from node `j` to node `i` .
 In the message-passing scheme, the incoming messages 
 from the neighborhood of `i` will later be aggregated
 in order to update (see [`update_node`](@ref)) the features of node `i`.
 
 The function operates on batches of edges, therefore
 `x_i`, `x_j`, and `e_ij` are tensors whose last dimension
-is the batch size. 
+is the batch size, or can be tuple/namedtuples of 
+such tensors, according to the input to propagate.
 
 By default, the function returns `x_j`.
 Custom layer should specialize this method with the desired behavior.
@@ -106,7 +108,7 @@ _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)
+function compute_batch_message(l, g::GNNGraph, x, e)
     s, t = edge_index(g)
     xi = _gather(x, t)
     xj = _gather(x, s)
@@ -121,12 +123,12 @@ _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)
+function aggregate_neighbors(g::GNNGraph, aggr, e)
     s, t = edge_index(g)
     _scatter(aggr, e, t)
 end
 
-aggregate_neighbors(l, g, aggr::Nothing, e) = nothing
+aggregate_neighbors(g::GNNGraph, aggr::Nothing, e) = nothing
 
 ## Step 3
 
@@ -137,6 +139,9 @@ Node update function for the GNN layer `l`,
 returning a new set of node features `x′` based on old 
 features `x` and the aggregated message `m̄` from the neighborhood.
 
+The input `m̄` is an array, a tuple or a named tuple, 
+reflecting the output of [`compute_message`](@ref).
+
 By default, the function returns `m̄`.
 Custom layers should  specialize this method with the desired behavior.
 
@@ -155,7 +160,7 @@ function update_node end
 Edge update function for the GNN layer `l`,
 returning a new set of edge features `e′` based on old 
 features `e` and the newly computed messages `m`
-from the [`message`](@ref) function.
+from the [`compute_message`](@ref) function.
 
 By default, the function returns `e`.
 Custom layers should specialize this method with the desired behavior.

test/cuda/msgpass.jl

@@ -10,19 +10,19 @@
         0 1 0 1 0 1
         0 1 1 0 1 0]
 
-    struct NewCudaLayer
+    struct NewCudaLayer{G} <: GNNLayer
         weight
     end
-    NewCudaLayer(m, n) = NewCudaLayer(randn(T, m, n))
-    @functor NewCudaLayer
+    NewCudaLayer{GRAPH_T}(m, n) = NewCudaLayer{GRAPH_T}(randn(T, m, n))
+    Flux.@functor NewCudaLayer{GRAPH_T}
 
-    (l::NewCudaLayer)(g, X) = GraphNeuralNetworks.propagate(l, g, +, X)
-    GraphNeuralNetworks.compute_message(n::NewCudaLayer, x_i, x_j, e_ij) = n.weight * x_j
-    GraphNeuralNetworks.update_node(::NewCudaLayer, m, x) = m
+    (l::NewCudaLayer{GRAPH_T})(g, X) = GraphNeuralNetworks.propagate(l, g, +, X)[1]
+    GraphNeuralNetworks.compute_message(n::NewCudaLayer{GRAPH_T}, x_i, x_j, e_ij) = n.weight * x_j
+    GraphNeuralNetworks.update_node(::NewCudaLayer{GRAPH_T}, m, x) = m
 
     X = rand(T, in_channel, N) |> gpu
     g = GNNGraph(adj, ndata=X, graph_type=GRAPH_T)
-    l = NewCudaLayer(out_channel, in_channel) |> gpu
+    l = NewCudaLayer{GRAPH_T}(out_channel, in_channel) |> gpu
 
     g_ = l(g)
     @test size(node_features(g_)) == (out_channel, N)

test/msgpass.jl

@@ -132,26 +132,25 @@ import GraphNeuralNetworks: compute_message, update_node, update_edge, propagate
             W
         end
 
-        NewLayerNT(in, out) = NewLayerW{GRAPH_T}(randn(T, out, in))
+        NewLayerNT(in, out) = NewLayerNT{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
+        function GraphNeuralNetworks.compute_message(l::NewLayerNT{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)
+        function GraphNeuralNetworks.update_node(l::NewLayerNT{GRAPH_T}, m, d) 
+            return (α=l.W * d.x + m.a + m.b, β=m)
         end
-        function update_edge(l::NewLayerW{GRAPH_T}, m, e) 
+        function GraphNeuralNetworks.update_edge(l::NewLayerNT{GRAPH_T}, m, e) 
             return m.a
         end
 
-        function (::NewLayerNT)(l, g, x, e)
-            x, e = propagate(l, g, (; x), (; e))
+        function (::NewLayerNT{GRAPH_T})(g, x, e)
+            x, e = propagate(l, g, mean, (; 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)