start implementing GNNChain

Author: CarloLucibello

Project.toml

@@ -11,6 +11,7 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

src/GraphNeuralNetworks.jl

@@ -9,6 +9,7 @@ import KrylovKit
 using CUDA
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
+using MacroTools: @forward
 using NNlib, NNlibCUDA
 using ChainRulesCore
 import LightGraphs

src/layers/basic.jl

@@ -6,3 +6,67 @@ An abstract type from which graph neural network layers are derived.
 See also [`GNNChain`](@ref).
 """
 abstract type GNNLayer end
+
+"""
+    GNNChain(layers...)
+    GNNChain(name = layer, ...)
+
+Collects multiple layers / functions to be called in sequence
+on a given input. Supports indexing and slicing, `m[2]` or `m[1:end-1]`,
+and if names are given, `m[:name] == m[1]` etc.
+
+## Examples
+
+```jldoctest
+julia> m = GNNChain(x -> x^2, x -> x+1);
+
+julia> m(5) == 26
+true
+
+julia> m = GNNChain(Dense(10, 5, tanh), Dense(5, 2));
+
+julia> x = rand(10, 32);
+
+julia> m(x) == m[2](m[1](x))
+true
+
+julia> m2 = GNNChain(enc = GNNChain(Flux.flatten, Dense(10, 5, tanh)), 
+                  dec = Dense(5, 2));
+
+                  julia> m2(x) == (m2[:dec] ∘ m2[:enc])(x)
+true
+```
+"""
+struct GNNChain{T}
+  layers::T
+
+  GNNChain(xs...) = new{typeof(xs)}(xs)
+  
+  function GNNChain(; kw...)
+    :layers in Base.keys(kw) && throw(ArgumentError("a GNNChain cannot have a named layer called `layers`"))
+    isempty(kw) && return new{Tuple{}}(())
+    new{typeof(values(kw))}(values(kw))
+  end
+end
+
+@forward GNNChain.layers Base.getindex, Base.length, Base.first, Base.last,
+    Base.iterate, Base.lastindex, Base.keys
+
+functor(::Type{<:GNNChain}, c) = c.layers, ls -> GNNChain(ls...)
+
+applychain(::Tuple{}, x) = x
+applychain(fs::Tuple, x) = applychain(tail(fs), first(fs)(x))
+
+(c::GNNChain)(x) = applychain(Tuple(c.layers), x)
+
+Base.getindex(c::GNNChain, i::AbstractArray) = GNNChain(c.layers[i]...)
+Base.getindex(c::GNNChain{<:NamedTuple}, i::AbstractArray) = 
+    GNNChain(; NamedTuple{Base.keys(c)[i]}(Tuple(c.layers)[i])...)
+
+function Base.show(io::IO, c::GNNChain)
+    print(io, "GNNChain(")
+    _show_layers(io, c.layers)
+    print(io, ")")
+end
+_show_layers(io, layers::Tuple) = join(io, layers, ", ")
+_show_layers(io, layers::NamedTuple) = join(io, ["$k = $v" for (k, v) in pairs(layers)], ", ")

src/layers/msgpass.jl

@@ -23,7 +23,7 @@ and [`message`](@ref) function, then call
 this method in the forward pass:
 
 ```julia
-function (l::GNNLayer)(g, X)
+function (l::MyLayer)(g, X)
     ... some prepocessing if needed ...
     E = nothing
     u = nothing