fix docs

Author: CarloLucibello

docs/src/messagepassing.md

@@ -20,13 +20,15 @@ In GraphNeuralNetworks.jl, the function [`propagate`](@ref) takes care of materi
 node features on each edge, applying the message function, performing the
 aggregation, and returning ``\bar{\mathbf{m}}``. 
 It is then left to the user to perform further node and edge updates,
-manypulating arrays of size ``D_{node} \times num\_nodes`` and   
+manipulating arrays of size ``D_{node} \times num\_nodes`` and   
 ``D_{edge} \times num\_edges``.
 
-As part of the [`propagate`](@ref) pipeline, we have the function
-[`apply_edges`](@ref). It can be independently used to materialize 
-node features on edges and perform edge-related computation without
-the following neighborhood aggregation one finds in `propagate`.
+[`propagate`](@ref) is composed of two steps corresponding to two
+exported methods:
+1.  [`apply_edges`](@ref) materializes node features on edges and 
+performs edge-related computation without. 
+2. [`aggregate_neighbors`](@ref) applies a reduction operator on the messages coming
+from the neighborhood of each node.
 
 The whole propagation mechanism internally relies on the [`NNlib.gather`](@ref) 
 and [`NNlib.scatter`](@ref) methods.
@@ -46,17 +48,11 @@ GNNGraph:
     num_nodes = 10
     num_edges = 20
 
+julia> x = ones(2,10);
 
-julia> x = ones(2,10)
-2×10 Matrix{Float64}:
- 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
- 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
-
-julia> z = 2ones(2,10)
-2×10 Matrix{Float64}:
- 2.0  2.0  2.0  2.0  2.0  2.0  2.0  2.0  2.0  2.0
- 2.0  2.0  2.0  2.0  2.0  2.0  2.0  2.0  2.0  2.0
+julia> z = 2ones(2,10);
 
+# Return an edge features arrays (D × num_edges)
 julia> apply_edges((xi, xj, e) -> xi .+ xj, g, xi=x, xj=z)
 2×20 Matrix{Float64}:
  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0  3.0
@@ -72,14 +68,16 @@ julia> apply_edges((xi, xj, e) -> xi.a + xi.b .* xj, g, xi=(a=x,b=z), xj=z)
  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0
  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0  5.0
 ```
+
 The function [`propagate`](@ref) instead performs the [`apply_edges`](@ref) operation
-but then also applies a reduction over each node's neighborhood.
+but then also applies a reduction over each node's neighborhood (see [`aggregate_neighbors`](@ref)).
 ```julia
 julia> propagate((xi, xj, e) -> xi .+ xj, g, +, xi=x, xj=z)
 2×10 Matrix{Float64}:
  3.0  6.0  9.0  9.0  0.0  6.0  6.0  3.0  15.0  3.0
  3.0  6.0  9.0  9.0  0.0  6.0  6.0  3.0  15.0  3.0
 
+# Previous output can be understood by looking at the degree
 julia> degree(g)
 10-element Vector{Int64}:
  1

src/layers/conv.jl

@@ -269,7 +269,7 @@ Implements the operation
 ```
 where the attention coefficients ``\alpha_{ij}`` are given by
 ```math
-\alpha_{ij} = \frac{1}{z_i} \exp(LeakyReLU(\mathbf{a}^T [W \mathbf{x}_i \,\|\, W \mathbf{x}_j]))
+\alpha_{ij} = \frac{1}{z_i} \exp(LeakyReLU(\mathbf{a}^T [W \mathbf{x}_i; W \mathbf{x}_j]))
 ```
 with ``z_i`` a normalization factor.
 
@@ -568,7 +568,7 @@ GraphSAGE convolution layer from paper [Inductive Representation Learning on Lar
 
 Performs:
 ```math
-\mathbf{x}_i' = W \cdot [\mathbf{x}_i \,\|\, \square_{j \in \mathcal{N}(i)} \mathbf{x}_j]
+\mathbf{x}_i' = W \cdot [\mathbf{x}_i; \square_{j \in \mathcal{N}(i)} \mathbf{x}_j]
 ```
 
 where the aggregation type is selected by `aggr`.
@@ -697,7 +697,7 @@ Performs the operation
 ```
 
 where ``\mathbf{z}_{ij}``  is the node and edge features concatenation 
-``[\mathbf{x}_i \| \mathbf{x}_j \| \mathbf{e}_{j\to i}]`` 
+``[\mathbf{x}_i; \mathbf{x}_j; \mathbf{e}_{j\to i}]`` 
 and ``\sigma`` is the sigmoid function.
 The residual ``\mathbf{x}_i`` is added only if `residual=true` and the output size is the same 
 as the input size.
@@ -829,12 +829,13 @@ end
 
 Convolution from [Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals](https://arxiv.org/pdf/1812.05055.pdf)
 paper. In the forward pass, takes as inputs node features `x` and edge features `e` and returns
-updated features `x̄, ē` according to 
+updated features `x'` and `e'` according to 
 
 ```math
-ē  = ϕe(vcat(xi, xj, e))
-x̄  = ϕv(vcat(x, \square_{j\in \mathcal{N}(i)} ē_{j\to i}))
+\mathbf{e}_{i\to j}'  = \phi_e([\mathbf{x}_i; \mathbf{x}_j; \mathbf{e}_{i\to j}])\\
+\mathbf{x}_{i}'  = \phi_v([\mathbf{x}_i; \square_{j\in \mathcal{N}(i)\,\mathbf{e}_{j\to i}'])
 ```
+
 `aggr` defines the aggregation to be performed.
 
 If the neural networks `ϕe` and  `ϕv` are not provided, they will be constructed from
@@ -849,7 +850,7 @@ g = rand_graph(10, 30)
 x = randn(3, 10)
 e = randn(3, 30)
 m = MEGNetConv(3 => 3)
-x̄, ē = m(g, x, e)
+x′, e′ = m(g, x, e)
 ```
 """
 struct MEGNetConv <: GNNLayer