docstring

Author: CarloLucibello

src/layers/basic.jl

@@ -12,29 +12,31 @@ abstract type GNNLayer end
     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
+on given input graph and input node features. 
 
-```
-julia> m = GNNChain(x -> x^2, x -> x+1);
+It allows to compose layers in a sequential fashion as `Flux.Chain`
+does, propagating the output of each layer to the next one.
+In addition, `GNNChain` handles the input graph as well, providing it 
+as a first argument only to layers subtyping the [`GNNLayer`](@ref) abstract type. 
 
-julia> m(5) == 26
-true
+`GNNChain` supports indexing and slicing, `m[2]` or `m[1:end-1]`,
+and if names are given, `m[:name] == m[1]` etc.
 
-julia> m = GNNChain(Dense(10, 5, tanh), Dense(5, 2));
+# Examples
 
-julia> x = rand(10, 32);
+```juliarepl
+julia> m = GNNChain(GCNConv(2=>5), BatchNorm(5), x -> relu.(x), Dense(5, 4));
 
-julia> m(x) == m[2](m[1](x))
-true
+julia> x = randn(Float32, 2, 3);
 
-julia> m2 = GNNChain(enc = GNNChain(Flux.flatten, Dense(10, 5, tanh)), 
-                  dec = Dense(5, 2));
+julia> g = GNNGraph([1,1,2,3], [2,3,1,1]);
 
-                  julia> m2(x) == (m2[:dec] ∘ m2[:enc])(x)
-true
+julia> m(g, x)
+4×3 Matrix{Float32}:
+  0.157941    0.15443     0.193471
+  0.0819516   0.0503105   0.122523
+  0.225933    0.267901    0.241878
+ -0.0134364  -0.0120716  -0.0172505
 ```
 """
 struct GNNChain{T}