models docs

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"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"

docs/make.jl

@@ -16,6 +16,7 @@ makedocs(;
                 "Basic Layers" => "api/basic.md",
                 "Convolutional Layers" => "api/conv.md",
                 "Pooling Layers" => "api/pool.md",
+                "NNlib" => "api/nnlib.md",
                ],
               "Developer Notes" => "dev.md",
             ],

docs/src/api/nnlib.md

@@ -0,0 +1,10 @@
+# NNlib
+
+Primitive functions implemented in NNlib.jl.
+
+```@docs
+NNlib.gather!
+NNlib.gather
+NNlib.scatter!
+NNlib.scatter
+```

docs/src/messagepassing.md

@@ -1,2 +1,7 @@
 # Message Passing
 
+```@docs
+message
+update
+propagate
+```

docs/src/models.md

@@ -1,30 +1,47 @@
 # Models
 
+GraphNeuralNetworks.jl provides common graph convolutional layers by which you can assemble arbitrarily deep or complex models. GNN layers are compatible with 
+Flux.jl ones, therefore expert Flux's users should be immediately able to define and train 
+their models. 
+
+In what follows, we discuss two different styles for model creation:
+the *explicit modeling* style, more verbose but more flexible, 
+and the *implicity modeling* style based on [`GNNChain`](@ref), more concise but less flexible.
+
 ## Explicit modeling
 
+In the explicit modeling style, the model is created according to the following steps:
+
+1. Define a new type for your model (`GNN` in the example below). Layers and submodels are fields.
+2. Apply `Flux.@functor` to the new type to make it Flux's compatible (parameters' collection, gpu movement, etc...)
+3. Optionally define a convenience constructor for your model.
+4. Define the forward pass by implementing the function call method for your type
+5. Instantiate the model. 
+
+Here is an example of this construction:
 ```julia
-using Flux, GraphNeuralNetworks
+using Flux, LightGraphs, GraphNeuralNetworks
 using Flux: @functor
 
-struct GNN
+struct GNN                                # step 1
     conv1
     bn
     conv2
     dropout
     dense
 end
 
-@functor GNN
+@functor GNN                              # step 2
 
-function GNN(din, d, dout)    
+function GNN(din::Int, d::Int, dout::Int) # step 3    
     GNN(GCNConv(din => d),
         BatchNorm(d),
         GraphConv(d => d, relu),
         Dropout(0.5),
         Dense(d, dout))
 end
 
-function (model::GNN)(g::GNNGraph, x)
+function (model::GNN)(g::GNNGraph, x)     # step 4
     x = model.conv1(g, x)
     x = relu.(model.bn(x))
     x = model.conv2(g, x)
@@ -34,19 +51,38 @@ function (model::GNN)(g::GNNGraph, x)
 end
 
 din, d, dout = 3, 4, 2 
-g = GNNGraph(random_regular_graph(10, 4), graph_type=GRAPH_T)
+g = GNNGraph(random_regular_graph(10, 4))
 X = randn(Float32, din, 10)
-model = GNN(din, d, dout)
+model = GNN(din, d, dout)                 # step 5
 y = model(g, X)
 ```
 
-## Compact modeling with GNNChains
+## Implicit modeling with GNNChains
+
+While very flexible, the way in which we defined `GNN` model definition in last section is a bit verbose.
+In order to simplify things, we provide the [`GNNChain`](@ref) type. It is very similar 
+to Flux's well known `Chain`. It allows to compose layers in a sequential fashion as Chain
+does, propagating the output of each layer to the next one. In addition, `GNNChain` 
+handles propagates the input graph as well, providing it as a first argument
+to layers subtyping the [`GNNLayer`](@ref) abstract type. 
+
+Using `GNNChain`, the previous example becomes
 
 ```julia
+using Flux, LightGraphs, GraphNeuralNetworks
+
+din, d, dout = 3, 4, 2 
+g = GNNGraph(random_regular_graph(10, 4))
+X = randn(Float32, din, 10)
+
 model = GNNChain(GCNConv(din => d),
                  BatchNorm(d),
                  x -> relu.(x),
                  GraphConv(d => d, relu),
                  Dropout(0.5),
                  Dense(d, dout))
+
+y = model(g, X)
 ```
+
+The `GNNChain` only propagates the graph and the node features. More complex scenarios, e.g. when also edge features are updated, have to be handled using the explicit definition of the forward pass. 

src/layers/conv.jl

@@ -334,6 +334,9 @@ end
 message(l::GatedGraphConv, x_i, x_j, e_ij) = x_j
 update(l::GatedGraphConv, m, x) = m
 
+# remove after https://github.com/JuliaDiff/ChainRules.jl/pull/521
+@non_differentiable fill!(x...)
+
 function (ggc::GatedGraphConv)(g::GNNGraph, H::AbstractMatrix{S}) where {T<:AbstractVector,S<:Real}
     check_num_nodes(g, H)
     m, n = size(H)