remove MessagePassing type, add GNNLayer

CarloLucibello · CarloLucibello · commit c2c6cfee15dd · 2021-08-31T13:06:13.000+02:00
diff --git a/src/GraphNeuralNetworks.jl b/src/GraphNeuralNetworks.jl
@@ -27,7 +27,10 @@ export
     adjacency_matrix, 
 
     # layers/msgpass
-    MessagePassing,
+
+    # layers/basic
+    GNNLayer,
+    GNNChain,
 
     # layers/conv
     GCNConv,
@@ -38,33 +41,23 @@ export
     EdgeConv,
     GINConv,
 
-    # layer/pool
+    # layers/pool
     GlobalPool,
     LocalPool,
     TopKPool,
-    topk_index,
+    topk_index
 
-    # models
-    GAE,
-    VGAE,
-    InnerProductDecoder,
-    VariationalEncoder,
-    summarize,
-    sample,
 
-    # layer/selector
-    bypass_graph
 
     
 include("gnngraph.jl")
 include("graph_conversions.jl")
 include("utils.jl")
 
 include("layers/msgpass.jl")
-
+include("layers/basic.jl")
 include("layers/conv.jl")
 include("layers/pool.jl")
-include("layers/misc.jl")
 
 
 end
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -0,0 +1,8 @@
+"""
+    abstract type GNNLayer end
+
+An abstract type from which graph neural network layers are derived.
+
+See also [`GNNChain`](@ref).
+"""
+abstract type GNNLayer end
diff --git a/src/layers/conv.jl b/src/layers/conv.jl
@@ -20,7 +20,7 @@ of size `(num_features, num_nodes)`.
 - `bias`: Add learnable bias.
 - `init`: Weights' initializer.
 """
-struct GCNConv{A<:AbstractMatrix, B, F} <: MessagePassing
+struct GCNConv{A<:AbstractMatrix, B, F} <: GNNLayer
     weight::A
     bias::B
     σ::F
@@ -97,7 +97,7 @@ with ``\hat{L}`` the [`scaled_laplacian`](@ref).
 - `bias`: Add learnable bias.
 - `init`: Weights' initializer.
 """
-struct ChebConv{A<:AbstractArray{<:Number,3}, B}
+struct ChebConv{A<:AbstractArray{<:Number,3}, B} <: GNNLayer
     weight::A
     bias::B
     k::Int
@@ -160,7 +160,7 @@ where the aggregation type is selected by `aggr`.
 - `bias`: Add learnable bias.
 - `init`: Weights' initializer.
 """
-struct GraphConv{A<:AbstractMatrix, B} <: MessagePassing
+struct GraphConv{A<:AbstractMatrix, B} <: GNNLayer
     weight1::A
     weight2::A
     bias::B
@@ -229,7 +229,7 @@ with ``z_i`` a normalization factor.
 - `concat`: Concatenate layer output or not. If not, layer output is averaged over the heads.
 - `negative_slope::Real`: Keyword argument, the parameter of LeakyReLU.
 """
-struct GATConv{T, A<:AbstractMatrix{T}, B} <: MessagePassing
+struct GATConv{T, A<:AbstractMatrix{T}, B} <: GNNLayer
     weight::A
     bias::B
     a::A
@@ -313,7 +313,7 @@ Implements the recursion
 - `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
 - `init`: Weight initialization function.
 """
-struct GatedGraphConv{A<:AbstractArray{<:Number,3}, R} <: MessagePassing
+struct GatedGraphConv{A<:AbstractArray{<:Number,3}, R} <: GNNLayer
     weight::A
     gru::R
     out_ch::Int
@@ -376,7 +376,7 @@ where `f` typically denotes a learnable function, e.g. a linear layer or a multi
 - `f`: A (possibly learnable) function acting on edge features. 
 - `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
 """
-struct EdgeConv <: MessagePassing
+struct EdgeConv <: GNNLayer
     nn
     aggr
 end
@@ -420,7 +420,7 @@ where `f` typically denotes a learnable function, e.g. a linear layer or a multi
 - `f`: A (possibly learnable) function acting on node features. 
 - `eps`: Weighting factor.
 """
-struct GINConv{R<:Real} <: MessagePassing
+struct GINConv{R<:Real} <: GNNLayer
     nn
     eps::R
 end
diff --git a/src/layers/misc.jl b/src/layers/misc.jl
diff --git a/src/layers/msgpass.jl b/src/layers/msgpass.jl
@@ -1,18 +1,9 @@
 # Adapted message passing from paper 
 # "Relational inductive biases, deep learning, and graph networks"
-"""
-    MessagePassing
-
-The abstract type from which all message passing layers are derived.
-
-Related methods are [`propagate`](@ref), [`message`](@ref),
-[`update`](@ref), [`update_edge`](@ref), and [`update_global`](@ref). 
-"""
-abstract type MessagePassing end
 
 """
-    propagate(mp::MessagePassing, g::GNNGraph, aggr)
-    propagate(mp::MessagePassing, g::GNNGraph, E, X, u, aggr)
+    propagate(mp, g::GNNGraph, aggr)
+    propagate(mp, g::GNNGraph, E, X, u, aggr)
 
 Perform the sequence of operation implementing the message-passing scheme
 and updating node, edge, and global features `X`, `E`, and `u` respectively.
@@ -27,9 +18,8 @@ X = update(mp, M̄, X, u)
 u = update_global(mp, E, X, u)
 ```
 
-Custom layers sub-typing [`MessagePassing`](@ref)
-typically call define their own [`update`](@ref)
-and [`message`](@ref) function, than call
+Custom layers typically define their own [`update`](@ref)
+and [`message`](@ref) function, then call
 this method in the forward pass:
 
 ```julia
@@ -45,14 +35,14 @@ See also [`message`](@ref) and [`update`](@ref).
 """
 function propagate end 
 
-function propagate(mp::MessagePassing, g::GNNGraph, aggr)
+function propagate(mp, g::GNNGraph, aggr)
     E, X, u = propagate(mp, g,
                         edge_feature(g), node_feature(g), global_feature(g), 
                         aggr)
     GNNGraph(g, nf=X, ef=E, gf=u)
 end
 
-function propagate(mp::MessagePassing, g::GNNGraph, E, X, u, aggr)
+function propagate(mp, g::GNNGraph, E, X, u, aggr)
     M = compute_batch_message(mp, g, E, X, u) 
     E = update_edge(mp, M, E, u)
     M̄ = aggregate_neighbors(mp, aggr, g, M)
@@ -62,7 +52,7 @@ function propagate(mp::MessagePassing, g::GNNGraph, E, X, u, aggr)
 end
 
 """
-    message(mp::MessagePassing, x_i, x_j, [e_ij, u])
+    message(mp, x_i, x_j, [e_ij, u])
 
 Message function for the message-passing scheme,
 returning the message from node `j` to node `i` .
@@ -71,12 +61,11 @@ from the neighborhood of `i` will later be aggregated
 in order to [`update`](@ref) the features of node `i`.
 
 By default, the function returns `x_j`.
-Layers subtyping [`MessagePassing`](@ref) should 
-specialize this method with custom behavior.
+Custom layer should specialize this method with the desired behavior.
 
 ## Arguments
 
-- `mp`: A [`MessagePassing`](@ref) layer.
+- `mp`: A gnn layer.
 - `x_i`: Features of the central node `i`.
 - `x_j`: Features of the neighbor `j` of node `i`.
 - `e_ij`: Features of edge (`i`, `j`).
@@ -87,20 +76,19 @@ See also [`update`](@ref) and [`propagate`](@ref).
 function message end 
 
 """
-    update(mp::MessagePassing, m̄, x, [u])
+    update(mp, m̄, x, [u])
 
 Update function for the message-passing scheme,
 returning a new set of node features `x′` based on old 
 features `x` and the incoming message from the neighborhood
 aggregation `m̄`.
 
 By default, the function returns `m̄`.
-Layers subtyping [`MessagePassing`](@ref) should 
-specialize this method with custom behavior.
+Custom layers should  specialize this method with the desired behavior.
 
 ## Arguments
 
-- `mp`: A [`MessagePassing`](@ref) layer.
+- `mp`: A gnn layer.
 - `m̄`: Aggregated edge messages from the [`message`](@ref) function.
 - `x`: Node features to be updated.
 - `u`: Global features.
@@ -115,41 +103,41 @@ _gather(x::Nothing, i) = nothing
 
 ## Step 1.
 
-function compute_batch_message(mp::MessagePassing, g, E, X, u)
+function compute_batch_message(mp, g, E, X, u)
     s, t = edge_index(g)
     Xi = _gather(X, t)
     Xj = _gather(X, s)
     M = message(mp, Xi, Xj, E, u)
     return M
 end
 
-# @inline message(mp::MessagePassing, i, j, x_i, x_j, e_ij, u) = message(mp, x_i, x_j, e_ij, u) # TODO add in the future
-@inline message(mp::MessagePassing, x_i, x_j, e_ij, u) = message(mp, x_i, x_j, e_ij)
-@inline message(mp::MessagePassing, x_i, x_j, e_ij) = message(mp, x_i, x_j)
-@inline message(mp::MessagePassing, x_i, x_j) = x_j
+# @inline message(mp, i, j, x_i, x_j, e_ij, u) = message(mp, x_i, x_j, e_ij, u) # TODO add in the future
+@inline message(mp, x_i, x_j, e_ij, u) = message(mp, x_i, x_j, e_ij)
+@inline message(mp, x_i, x_j, e_ij) = message(mp, x_i, x_j)
+@inline message(mp, x_i, x_j) = x_j
 
 ## Step 2
 
-@inline update_edge(mp::MessagePassing, M, E, u) = update_edge(mp, M, E)
-@inline update_edge(mp::MessagePassing, M, E) = E
+@inline update_edge(mp, M, E, u) = update_edge(mp, M, E)
+@inline update_edge(mp, M, E) = E
 
 ##  Step 3
 
-function aggregate_neighbors(mp::MessagePassing, aggr, g, E)
+function aggregate_neighbors(mp, aggr, g, E)
     s, t = edge_index(g)
     NNlib.scatter(aggr, E, t)
 end
 
-aggregate_neighbors(mp::MessagePassing, aggr::Nothing, g, E) = nothing
+aggregate_neighbors(mp, aggr::Nothing, g, E) = nothing
 
 ## Step 4
 
-# @inline update(mp::MessagePassing, i, m̄, x, u) = update(mp, m, x, u)
-@inline update(mp::MessagePassing, m̄, x, u) = update(mp, m̄, x)
-@inline update(mp::MessagePassing, m̄, x) = m̄
+# @inline update(mp, i, m̄, x, u) = update(mp, m, x, u)
+@inline update(mp, m̄, x, u) = update(mp, m̄, x)
+@inline update(mp, m̄, x) = m̄
 
 ## Step 5
 
-@inline update_global(mp::MessagePassing, E, X, u) = u
+@inline update_global(mp, E, X, u) = u
 
 ### end steps ###
diff --git a/test/cuda/layers/msgpass.jl b/test/cuda/layers/msgpass.jl
@@ -9,7 +9,7 @@ adj = [0 1 0 0 0 0
        0 1 0 1 0 1
        0 1 1 0 1 0]
 
-struct NewCudaLayer <: MessagePassing
+struct NewCudaLayer
     weight
 end
 NewCudaLayer(m, n) = NewCudaLayer(randn(T, m,n))
diff --git a/test/layers/basic.jl b/test/layers/basic.jl
@@ -0,0 +1,3 @@
+@testset "basic" begin
+
+end
diff --git a/test/layers/misc.jl b/test/layers/misc.jl
diff --git a/test/layers/msgpass.jl b/test/layers/msgpass.jl
@@ -1,4 +1,4 @@
-@testset "MessagePassing" begin 
+@testset "message passing" begin 
     in_channel = 10
     out_channel = 5
     num_V = 6
@@ -12,7 +12,7 @@
             0 1 0 1 0 1
             0 1 1 0 1 0]
 
-    struct NewLayer{G} <: MessagePassing end
+    struct NewLayer{G} end
 
     X = rand(T, in_channel, num_V)
     E = rand(T, in_channel, num_E)
@@ -90,7 +90,7 @@
         @test size(global_feature(fg_)) == (in_channel,)
     end
 
-    struct NewLayerW{G} <: MessagePassing
+    struct NewLayerW{G}
         weight
     end
 
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -16,9 +16,9 @@ include("cuda/test_utils.jl")
 tests = [
     "gnngraph",
     "layers/msgpass",
+    "layers/basic",
     "layers/conv",
     "layers/pool",
-    "layers/misc",
 ]
 
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+@testset "basic" begin`
	`2`	`+`
	`3`	`+end`