From 7e326155ef258ffbbec0266a1762662d1b5887cd Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 5 Feb 2025 23:36:45 +0100
Subject: [PATCH 1/8] fix doctests

---
 GNNGraphs/docs/make.jl                         |  1 -
 GNNGraphs/docs/src/guides/heterograph.md       | 10 +++++-----
 GNNGraphs/docs/src/guides/temporalgraph.md     |  6 +++++-
 GNNGraphs/src/gnnheterograph/gnnheterograph.jl |  6 +++---
 GNNGraphs/src/mldatasets.jl                    | 10 +++++-----
 GNNGraphs/src/transform.jl                     | 10 +++++-----
 GNNLux/docs/make.jl                            |  1 -
 GNNlib/docs/make.jl                            |  1 -
 GraphNeuralNetworks/docs/make.jl               |  1 -
 9 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/GNNGraphs/docs/make.jl b/GNNGraphs/docs/make.jl
index 97ccb2060..a6daad102 100644
--- a/GNNGraphs/docs/make.jl
+++ b/GNNGraphs/docs/make.jl
@@ -25,7 +25,6 @@ mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/ma
 
 makedocs(;
     modules = [GNNGraphs],
-    doctest = false, # TODO enable doctest
     format = Documenter.HTML(; mathengine, 
                     prettyurls = get(ENV, "CI", nothing) == "true", 
                     assets = [],
diff --git a/GNNGraphs/docs/src/guides/heterograph.md b/GNNGraphs/docs/src/guides/heterograph.md
index 525a64474..c24d7391c 100644
--- a/GNNGraphs/docs/src/guides/heterograph.md
+++ b/GNNGraphs/docs/src/guides/heterograph.md
@@ -88,13 +88,11 @@ julia> g.etypes  # edge types
 
 Node, edge, and graph features can be added at construction time or later using:
 ```jldoctest hetero
-# equivalent to g.ndata[:user][:x] = ...
-julia> g[:user].x = rand(Float32, 64, 3);
+julia> g[:user].x = rand(Float32, 64, 3); # equivalent to g.ndata[:user][:x] = ...
 
 julia> g[:movie].z = rand(Float32, 64, 13);
 
-# equivalent to g.edata[(:user, :rate, :movie)][:e] = ...
-julia> g[:user, :rate, :movie].e = rand(Float32, 64, 4);
+julia> g[:user, :rate, :movie].e = rand(Float32, 64, 4); # equivalent to g.edata[(:user, :rate, :movie)][:e] = ...
 
 julia> g
 GNNHeteroGraph:
@@ -110,6 +108,8 @@ GNNHeteroGraph:
 ## Batching
 Similarly to graphs, also heterographs can be batched together.
 ```jldoctest hetero
+julia> using MLUtils
+
 julia> gs = [rand_bipartite_heterograph((5, 10), 20) for _ in 1:32];
 
 julia> MLUtils.batch(gs)
@@ -140,4 +140,4 @@ end
 
 ## Graph convolutions on heterographs
 
-See `HeteroGraphConv` for how to perform convolutions on heterogeneous graphs.
+See [`HeteroGraphConv`](@ref) for how to perform convolutions on heterogeneous graphs.
diff --git a/GNNGraphs/docs/src/guides/temporalgraph.md b/GNNGraphs/docs/src/guides/temporalgraph.md
index 631e3b944..4458e87fa 100644
--- a/GNNGraphs/docs/src/guides/temporalgraph.md
+++ b/GNNGraphs/docs/src/guides/temporalgraph.md
@@ -65,6 +65,10 @@ Snapshots in a temporal graph can be accessed using indexing:
 julia> snapshots = [rand_graph(10, 20), rand_graph(10, 14), rand_graph(10, 22)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
+TemporalSnapshotsGNNGraph:
+  num_nodes: [10, 10, 10]
+  num_edges: [20, 14, 22]
+  num_snapshots: 3
 
 julia> tg[1] # first snapshot
 GNNGraph:
@@ -169,7 +173,7 @@ TemporalSnapshotsGNNGraph:
   num_edges: [20, 14, 22]
   num_snapshots: 3
   tgdata:
-        y = 3×1 Matrix{Float32}
+    y = 3×1 Matrix{Float32}
 
 julia> tg.ndata # vector of DataStore containing node features for each snapshot
 3-element Vector{DataStore}:
diff --git a/GNNGraphs/src/gnnheterograph/gnnheterograph.jl b/GNNGraphs/src/gnnheterograph/gnnheterograph.jl
index 6594ca55f..4a3f8b924 100644
--- a/GNNGraphs/src/gnnheterograph/gnnheterograph.jl
+++ b/GNNGraphs/src/gnnheterograph/gnnheterograph.jl
@@ -198,18 +198,18 @@ function Base.show(io::IO, ::MIME"text/plain", g::GNNHeteroGraph)
             print(io, "\n  ndata:")
             for k in sort(collect(keys(g.ndata)))
                 isempty(g.ndata[k]) && continue    
-                print(io, "\n\t", _str(k), "  =>  $(shortsummary(g.ndata[k]))")
+                print(io, "\n    ", _str(k), "  =>  $(shortsummary(g.ndata[k]))")
             end
         end
         if !isempty(g.edata) && !all(isempty, values(g.edata))
             print(io, "\n  edata:")
             for k in sort(collect(keys(g.edata)))
                 isempty(g.edata[k]) && continue
-                print(io, "\n\t$k  =>  $(shortsummary(g.edata[k]))")
+                print(io, "\n    $k  =>  $(shortsummary(g.edata[k]))")
             end
         end
         if !isempty(g.gdata)
-            print(io, "\n  gdata:\n\t")
+            print(io, "\n  gdata:\n    ")
             shortsummary(io, g.gdata)
         end
     end
diff --git a/GNNGraphs/src/mldatasets.jl b/GNNGraphs/src/mldatasets.jl
index 40568a6c7..fe0ca94fe 100644
--- a/GNNGraphs/src/mldatasets.jl
+++ b/GNNGraphs/src/mldatasets.jl
@@ -15,11 +15,11 @@ GNNGraph:
   num_nodes: 2708
   num_edges: 10556
   ndata:
-        val_mask = 2708-element BitVector
-        targets = 2708-element Vector{Int64}
-        test_mask = 2708-element BitVector
-        features = 1433×2708 Matrix{Float32}
-        train_mask = 2708-element BitVector
+    val_mask = 2708-element BitVector
+    targets = 2708-element Vector{Int64}
+    test_mask = 2708-element BitVector
+    features = 1433×2708 Matrix{Float32}
+    train_mask = 2708-element BitVector
 ```
 """
 function mldataset2gnngraph(dataset::D) where {D}
diff --git a/GNNGraphs/src/transform.jl b/GNNGraphs/src/transform.jl
index ea1fe9583..afc94b517 100644
--- a/GNNGraphs/src/transform.jl
+++ b/GNNGraphs/src/transform.jl
@@ -459,14 +459,14 @@ GNNGraph:
   num_nodes: 4
   num_edges: 5
   edata:
-        e = 5-element Vector{Float64}
+    e = 5-element Vector{Float64}
 
 julia> g2 = to_bidirected(g)
 GNNGraph:
   num_nodes: 4
   num_edges: 7
   edata:
-        e = 7-element Vector{Float64}
+    e = 7-element Vector{Float64}
 
 julia> edge_index(g2)
 ([1, 2, 2, 3, 3, 4, 4], [2, 1, 3, 2, 4, 3, 4])
@@ -644,14 +644,14 @@ GNNGraph:
   num_nodes: 4
   num_edges: 4
   ndata:
-        x = 3×4 Matrix{Float32}
+    x = 3×4 Matrix{Float32}
 
 julia> g2 = rand_graph(5, 4, ndata=zeros(Float32, 3, 5))
 GNNGraph:
   num_nodes: 5
   num_edges: 4
   ndata:
-        x = 3×5 Matrix{Float32}
+    x = 3×5 Matrix{Float32}
 
 julia> g12 = MLUtils.batch([g1, g2])
 GNNGraph:
@@ -659,7 +659,7 @@ GNNGraph:
   num_edges: 8
   num_graphs: 2
   ndata:
-        x = 3×9 Matrix{Float32}
+    x = 3×9 Matrix{Float32}
 
 julia> g12.ndata.x
 3×9 Matrix{Float32}:
diff --git a/GNNLux/docs/make.jl b/GNNLux/docs/make.jl
index aa4fedf91..778d8200f 100644
--- a/GNNLux/docs/make.jl
+++ b/GNNLux/docs/make.jl
@@ -37,7 +37,6 @@ cp(joinpath(@__DIR__, "../../GNNlib/docs/src"),
 
 makedocs(;
     modules = [GNNLux, GNNGraphs, GNNlib],
-    doctest = false, # TODO: enable doctest
     plugins = [interlinks],
     format = Documenter.HTML(; mathengine, 
                     prettyurls = get(ENV, "CI", nothing) == "true", 
diff --git a/GNNlib/docs/make.jl b/GNNlib/docs/make.jl
index 9d1cc097b..37f0aa533 100644
--- a/GNNlib/docs/make.jl
+++ b/GNNlib/docs/make.jl
@@ -26,7 +26,6 @@ cp(joinpath(@__DIR__, "../../GNNGraphs/docs/src/"),
 
 makedocs(;
     modules = [GNNlib, GNNGraphs],
-    doctest = false, # TODO enable doctest
     plugins = [interlinks],
     format = Documenter.HTML(; mathengine, prettyurls, assets = assets, size_threshold=nothing),
     sitename = "GNNlib.jl",
diff --git a/GraphNeuralNetworks/docs/make.jl b/GraphNeuralNetworks/docs/make.jl
index 94ada2d39..071403e19 100644
--- a/GraphNeuralNetworks/docs/make.jl
+++ b/GraphNeuralNetworks/docs/make.jl
@@ -37,7 +37,6 @@ cp(joinpath(@__DIR__, "../../GNNlib/docs/src"),
 
 makedocs(;
     modules = [GraphNeuralNetworks, GNNGraphs, GNNlib],
-    doctest = false, # TODO: enable doctest
     plugins = [interlinks],
     format = Documenter.HTML(; mathengine, 
                             prettyurls = get(ENV, "CI", nothing) == "true", 

From 50d51eff6673c62080be213d97b001682c88bfa6 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 00:46:39 +0100
Subject: [PATCH 2/8] fix gnngraphs doctests

---
 GNNGraphs/docs/src/guides/heterograph.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/GNNGraphs/docs/src/guides/heterograph.md b/GNNGraphs/docs/src/guides/heterograph.md
index c24d7391c..2859eee9a 100644
--- a/GNNGraphs/docs/src/guides/heterograph.md
+++ b/GNNGraphs/docs/src/guides/heterograph.md
@@ -99,10 +99,10 @@ GNNHeteroGraph:
   num_nodes: Dict(:movie => 13, :user => 3)
   num_edges: Dict((:user, :rate, :movie) => 4)
   ndata:
-        :movie  =>  DataStore(z = [64×13 Matrix{Float32}])
-        :user  =>  DataStore(x = [64×3 Matrix{Float32}])
+    :movie  =>  DataStore(z = [64×13 Matrix{Float32}])
+    :user  =>  DataStore(x = [64×3 Matrix{Float32}])
   edata:
-        (:user, :rate, :movie)  =>  DataStore(e = [64×4 Matrix{Float32}])
+    (:user, :rate, :movie)  =>  DataStore(e = [64×4 Matrix{Float32}])
 ```
 
 ## Batching
@@ -140,4 +140,4 @@ end
 
 ## Graph convolutions on heterographs
 
-See [`HeteroGraphConv`](@ref) for how to perform convolutions on heterogeneous graphs.
+See `HeteroGraphConv` for how to perform convolutions on heterogeneous graphs.

From d781330324355f5b0b68266244adf878d6f11300 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 01:17:09 +0100
Subject: [PATCH 3/8] fix GraphNeuralNetworks.jl

---
 GraphNeuralNetworks/docs/make.jl              |  1 +
 GraphNeuralNetworks/src/layers/basic.jl       | 33 +++++++-----------
 GraphNeuralNetworks/src/layers/heteroconv.jl  |  6 ++--
 .../src/layers/temporalconv.jl                | 34 ++++++++++++-------
 4 files changed, 38 insertions(+), 36 deletions(-)

diff --git a/GraphNeuralNetworks/docs/make.jl b/GraphNeuralNetworks/docs/make.jl
index 071403e19..13b6bfa91 100644
--- a/GraphNeuralNetworks/docs/make.jl
+++ b/GraphNeuralNetworks/docs/make.jl
@@ -12,6 +12,7 @@ using GraphNeuralNetworks
 using Flux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
 
+DocMeta.setdocmeta!(GNNGraphs, :DocTestSetup, :(using GNNGraphs, MLUtils); recursive = true)
 DocMeta.setdocmeta!(GraphNeuralNetworks, :DocTestSetup, :(using GraphNeuralNetworks); recursive = true)
 
 mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/mathtools"]),
diff --git a/GraphNeuralNetworks/src/layers/basic.jl b/GraphNeuralNetworks/src/layers/basic.jl
index 4f99ddba4..45a1e0b05 100644
--- a/GraphNeuralNetworks/src/layers/basic.jl
+++ b/GraphNeuralNetworks/src/layers/basic.jl
@@ -74,30 +74,22 @@ julia> using Flux, GraphNeuralNetworks
 julia> m = GNNChain(GCNConv(2=>5), 
                     BatchNorm(5), 
                     x -> relu.(x), 
-                    Dense(5, 4))
-GNNChain(GCNConv(2 => 5), BatchNorm(5), #7, Dense(5 => 4))
+                    Dense(5, 4));
 
 julia> x = randn(Float32, 2, 3);
 
 julia> g = rand_graph(3, 6)
 GNNGraph:
-    num_nodes = 3
-    num_edges = 6
+  num_nodes: 3
+  num_edges: 6
 
-julia> m(g, x)
-4×3 Matrix{Float32}:
-    -0.795592  -0.795592  -0.795592
-    -0.736409  -0.736409  -0.736409
-    0.994925   0.994925   0.994925
-    0.857549   0.857549   0.857549
+julia> m(g, x) |> size
+(4, 3)
 
-julia> m2 = GNNChain(enc = m, 
-                     dec = DotDecoder())
-GNNChain(enc = GNNChain(GCNConv(2 => 5), BatchNorm(5), #7, Dense(5 => 4)), dec = DotDecoder())
+julia> m2 = GNNChain(enc = m, dec = DotDecoder());
 
-julia> m2(g, x)
-1×6 Matrix{Float32}:
- 2.90053  2.90053  2.90053  2.90053  2.90053  2.90053
+julia> m2(g, x) |> size
+(1, 6)
 
 julia> m2[:enc](g, x) == m(g, x)
 true
@@ -196,15 +188,14 @@ returns the dot product `x_i ⋅ xj` on each edge.
 ```jldoctest
 julia> g = rand_graph(5, 6)
 GNNGraph:
-    num_nodes = 5
-    num_edges = 6
+  num_nodes: 5
+  num_edges: 6
 
 julia> dotdec = DotDecoder()
 DotDecoder()
 
-julia> dotdec(g, rand(2, 5))
-1×6 Matrix{Float64}:
- 0.345098  0.458305  0.106353  0.345098  0.458305  0.106353
+julia> dotdec(g, rand(2, 5)) |> size
+(1, 6)
 ```
 """
 struct DotDecoder <: GNNLayer end
diff --git a/GraphNeuralNetworks/src/layers/heteroconv.jl b/GraphNeuralNetworks/src/layers/heteroconv.jl
index a10ebb0c7..7666ac8e8 100644
--- a/GraphNeuralNetworks/src/layers/heteroconv.jl
+++ b/GraphNeuralNetworks/src/layers/heteroconv.jl
@@ -21,10 +21,12 @@ have to be aggregated using the `aggr` function. The default is to sum the outpu
 # Examples 
 
 ```jldoctest
-julia> g = rand_bipartite_heterograph((10, 15), 20)
+julia> using GraphNeuralNetworks, Flux
+
+julia> g = rand_bipartite_heterograph((10, 15), 80)
 GNNHeteroGraph:
   num_nodes: Dict(:A => 10, :B => 15)
-  num_edges: Dict((:A, :to, :B) => 20, (:B, :to, :A) => 20)
+  num_edges: Dict((:A, :to, :B) => 80, (:B, :to, :A) => 80)
 
 julia> x = (A = rand(Float32, 64, 10), B = rand(Float32, 64, 15));
 
diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index bc12f2fab..2872529d6 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -59,16 +59,13 @@ Returns the updated node features:
 The following example considers a static graph and a time-varying node features.
 
 ```jldoctest
-julia> num_nodes, num_edges = 5, 10;
+julia> num_nodes, num_edges = 5, 16;
 
 julia> d_in, d_out = 2, 3;
 
 julia> timesteps = 5;
 
 julia> g = rand_graph(num_nodes, num_edges);
-GNNGraph:
-  num_nodes: 5
-  num_edges: 10
 
 julia> x = rand(Float32, d_in, timesteps, num_nodes);
 
@@ -93,11 +90,15 @@ julia> timesteps = 5;
 
 julia> num_nodes = [10, 10, 10, 10, 10];
 
-julia> num_edges = [10, 12, 14, 16, 18];
+julia> num_edges = [20, 22, 24, 26, 28];
 
 julia> snapshots = [rand_graph(n, m) for (n, m) in zip(num_nodes, num_edges)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
+TemporalSnapshotsGNNGraph:
+  num_nodes: [10, 10, 10, 10, 10]
+  num_edges: [20, 22, 24, 26, 28]
+  num_snapshots: 5
 
 julia> x = [rand(Float32, d_in, n) for n in num_nodes];
 
@@ -269,7 +270,7 @@ See [`GNNRecurrence`](@ref) for more details.
 # Examples
 
 ```jldoctest
-julia> num_nodes, num_edges = 5, 10;
+julia> num_nodes, num_edges = 5, 16;
 
 julia> d_in, d_out = 2, 3;
 
@@ -280,7 +281,7 @@ julia> g = rand_graph(num_nodes, num_edges);
 julia> x = rand(Float32, d_in, timesteps, num_nodes);
 
 julia> layer = GConvGRU(d_in => d_out, 2)
-GConvGRU(
+GNNRecurrence(
   GConvGRUCell(2 => 3, 2),              # 108 parameters
 )                   # Total: 12 arrays, 108 parameters, 1.148 KiB.
 
@@ -326,9 +327,9 @@ where `output` is the updated hidden state `h` of the LSTM cell and `state` is t
 # Examples
 
 ```jldoctest
-julia> using GraphNeuralNetworks, Flux
+julia> using Flux
 
-julia> num_nodes, num_edges = 5, 10;
+julia> num_nodes, num_edges = 5, 16;
 
 julia> d_in, d_out = 2, 3;
 
@@ -453,7 +454,7 @@ See [`GNNRecurrence`](@ref) for more details.
 # Examples
 
 ```jldoctest
-julia> num_nodes, num_edges = 5, 10;
+julia> num_nodes, num_edges = 5, 16;
 
 julia> d_in, d_out = 2, 3;
 
@@ -727,15 +728,19 @@ julia> timesteps = 5;
 
 julia> num_nodes = [10, 10, 10, 10, 10];
 
-julia> num_edges = [10, 12, 14, 16, 18];
+julia> num_edges = [60, 62, 64, 66, 68];
 
 julia> snapshots = [rand_graph(n, m) for (n, m) in zip(num_nodes, num_edges)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
+TemporalSnapshotsGNNGraph:
+  num_nodes: [10, 10, 10, 10, 10]
+  num_edges: [60, 62, 64, 66, 68]
+  num_snapshots: 5
 
 julia> x = [rand(Float32, d_in, n) for n in num_nodes];
 
-julia> cell = EvolveGCNO(d_in => d_out)
+julia> layer = EvolveGCNO(d_in => d_out)
 GNNRecurrence(
   EvolveGCNOCell(2 => 3),               # 321 parameters
 )                   # Total: 5 arrays, 321 parameters, 1.535 KiB.
@@ -743,7 +748,7 @@ GNNRecurrence(
 julia> y = layer(tg, x);
 
 julia> length(y)    # timesteps
-5 
+5
 
 julia> size(y[end]) # (d_out, num_nodes[end])
 (3, 10)
@@ -874,6 +879,9 @@ julia> g = rand_graph(num_nodes, num_edges);
 julia> x = rand(Float32, d_in, timesteps, num_nodes);
 
 julia> layer = TGCN(d_in => d_out)
+GNNRecurrence(
+  TGCNCell(2 => 3),                     # 126 parameters
+)                   # Total: 18 arrays, 126 parameters, 1.469 KiB.
 
 julia> y = layer(g, x);
 

From 86fad04dcd6e8e43ca2581206ac9f91af5cd3582 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 01:20:59 +0100
Subject: [PATCH 4/8] datadeps

---
 GNNGraphs/docs/make.jl           | 1 +
 GraphNeuralNetworks/docs/make.jl | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/GNNGraphs/docs/make.jl b/GNNGraphs/docs/make.jl
index a6daad102..ce02a7843 100644
--- a/GNNGraphs/docs/make.jl
+++ b/GNNGraphs/docs/make.jl
@@ -10,6 +10,7 @@ using MLUtils # this is needed by setdocmeta!
 import Graphs
 using Graphs: induced_subgraph
 
+ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
 
 DocMeta.setdocmeta!(GNNGraphs, :DocTestSetup, :(using GNNGraphs, MLUtils); recursive = true)
 
diff --git a/GraphNeuralNetworks/docs/make.jl b/GraphNeuralNetworks/docs/make.jl
index 13b6bfa91..7090f479e 100644
--- a/GraphNeuralNetworks/docs/make.jl
+++ b/GraphNeuralNetworks/docs/make.jl
@@ -12,6 +12,8 @@ using GraphNeuralNetworks
 using Flux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
 
+ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
+
 DocMeta.setdocmeta!(GNNGraphs, :DocTestSetup, :(using GNNGraphs, MLUtils); recursive = true)
 DocMeta.setdocmeta!(GraphNeuralNetworks, :DocTestSetup, :(using GraphNeuralNetworks); recursive = true)
 

From ed569789b65adba1028a16c05537f567f254a320 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 07:40:49 +0100
Subject: [PATCH 5/8] fix doctests

---
 GNNLux/docs/make.jl                      |  4 ++++
 GNNLux/docs/src/api/pool.md              |  7 -------
 GNNLux/src/layers/basic.jl               | 17 ++++++++++++-----
 GNNlib/docs/make.jl                      |  4 ++++
 GraphNeuralNetworks/docs/src/api/pool.md |  7 -------
 5 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/GNNLux/docs/make.jl b/GNNLux/docs/make.jl
index 778d8200f..08c5d89da 100644
--- a/GNNLux/docs/make.jl
+++ b/GNNLux/docs/make.jl
@@ -12,6 +12,10 @@ using GNNLux
 using Lux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
 
+ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
+
+DocMeta.setdocmeta!(GNNGraphs, :DocTestSetup, :(using GNNGraphs, MLUtils); recursive = true)
+DocMeta.setdocmeta!(GNNlib, :DocTestSetup, :(using GNNlib); recursive = true)
 DocMeta.setdocmeta!(GNNLux, :DocTestSetup, :(using GNNLux); recursive = true)
 
 mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/mathtools"]),
diff --git a/GNNLux/docs/src/api/pool.md b/GNNLux/docs/src/api/pool.md
index a6d6d7f8b..42f9f67bc 100644
--- a/GNNLux/docs/src/api/pool.md
+++ b/GNNLux/docs/src/api/pool.md
@@ -5,13 +5,6 @@ CollapsedDocStrings = true
 
 # Pooling Layers
 
-## Index
-
-```@index
-Order = [:type, :function]
-Pages   = ["pool.md"]
-```
-
 ```@autodocs
 Modules = [GNNLux]
 Pages   = ["layers/pool.jl"]
diff --git a/GNNLux/src/layers/basic.jl b/GNNLux/src/layers/basic.jl
index ac28cabb7..ea0df2908 100644
--- a/GNNLux/src/layers/basic.jl
+++ b/GNNLux/src/layers/basic.jl
@@ -31,9 +31,14 @@ julia> using Lux, GNNLux, Random
 
 julia> rng = Random.default_rng();
 
-julia> m = GNNChain(GCNConv(2=>5), 
-                    x -> relu.(x), 
-                    Dense(5=>4))
+julia> m = GNNChain(GCNConv(2 => 5, relu), Dense(5 => 4))
+GNNChain(
+    layers = NamedTuple(
+        layer_1 = GCNConv(2 => 5, relu),  # 15 parameters
+        layer_2 = Dense(5 => 4),        # 24 parameters
+    ),
+)         # Total: 39 parameters,
+          #        plus 0 states.
 
 julia> x = randn(rng, Float32, 2, 3);
 
@@ -44,8 +49,10 @@ GNNGraph:
 
 julia> ps, st = LuxCore.setup(rng, m);
 
-julia> m(g, x, ps, st)     # First entry is the output, second entry is the state of the model
-(Float32[-0.15594329 -0.15594329 -0.15594329; 0.93431795 0.93431795 0.93431795; 0.27568763 0.27568763 0.27568763; 0.12568939 0.12568939 0.12568939], (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()))
+julia> y, st = m(g, x, ps, st);     # First entry is the output, second entry is the state of the model
+
+julia> size(y)
+(4, 3)
 ```
 """
 @concrete struct GNNChain <: GNNContainerLayer{(:layers,)}
diff --git a/GNNlib/docs/make.jl b/GNNlib/docs/make.jl
index 37f0aa533..7538b2bd3 100644
--- a/GNNlib/docs/make.jl
+++ b/GNNlib/docs/make.jl
@@ -12,6 +12,10 @@ using GNNGraphs
 import Graphs
 using DocumenterInterLinks
 
+ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
+DocMeta.setdocmeta!(GNNGraphs, :DocTestSetup, :(using GNNGraphs, MLUtils); recursive = true)
+DocMeta.setdocmeta!(GNNlib, :DocTestSetup, :(using GNNlib); recursive = true)
+
 assets=[]
 prettyurls = get(ENV, "CI", nothing) == "true"
 mathengine = MathJax3()
diff --git a/GraphNeuralNetworks/docs/src/api/pool.md b/GraphNeuralNetworks/docs/src/api/pool.md
index 52e7cbe9b..c0cd17280 100644
--- a/GraphNeuralNetworks/docs/src/api/pool.md
+++ b/GraphNeuralNetworks/docs/src/api/pool.md
@@ -5,13 +5,6 @@ CollapsedDocStrings = true
 
 # Pooling Layers
 
-## Index
-
-```@index
-Order = [:type, :function]
-Pages   = ["pool.md"]
-```
-
 ```@autodocs
 Modules = [GraphNeuralNetworks]
 Pages   = ["layers/pool.jl"]

From 43defacb397c209c527473f08ac836f9cef1b8de Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 07:49:37 +0100
Subject: [PATCH 6/8] cleanup

---
 GNNGraphs/docs/src/api/heterograph.md | 1 +
 GNNLux/docs/make.jl                   | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/GNNGraphs/docs/src/api/heterograph.md b/GNNGraphs/docs/src/api/heterograph.md
index 18a5562e2..5789715fa 100644
--- a/GNNGraphs/docs/src/api/heterograph.md
+++ b/GNNGraphs/docs/src/api/heterograph.md
@@ -38,3 +38,4 @@ Modules = [GNNGraphs]
 Pages   = ["gnnheterograph/generate.jl"]
 Private = false
 ```
+å
\ No newline at end of file
diff --git a/GNNLux/docs/make.jl b/GNNLux/docs/make.jl
index 08c5d89da..d79dd6171 100644
--- a/GNNLux/docs/make.jl
+++ b/GNNLux/docs/make.jl
@@ -85,7 +85,7 @@ makedocs(;
             "Layers" => [
                 "Basic layers" => "api/basic.md",
                 "Convolutional layers" => "api/conv.md",
-                # "Pooling layers" => "api/pool.md",
+                "Pooling layers" => "api/pool.md",
                 "Temporal Convolutional layers" => "api/temporalconv.md",
                 # "Hetero Convolutional layers" => "api/heteroconv.md",
             ]

From 92c1b9a4ab247c8a1de917754115f78265827d2d Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 07:56:05 +0100
Subject: [PATCH 7/8] add mlutils

---
 GNNLux/docs/Project.toml              | 1 +
 GNNLux/docs/make.jl                   | 1 +
 GNNlib/docs/Project.toml              | 1 +
 GraphNeuralNetworks/docs/Project.toml | 1 +
 GraphNeuralNetworks/docs/make.jl      | 2 ++
 5 files changed, 6 insertions(+)

diff --git a/GNNLux/docs/Project.toml b/GNNLux/docs/Project.toml
index 36822253f..206899f94 100644
--- a/GNNLux/docs/Project.toml
+++ b/GNNLux/docs/Project.toml
@@ -11,6 +11,7 @@ Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 TSne = "24678dba-d5e9-5843-a4c6-250288b04835"
diff --git a/GNNLux/docs/make.jl b/GNNLux/docs/make.jl
index d79dd6171..3c39c97f4 100644
--- a/GNNLux/docs/make.jl
+++ b/GNNLux/docs/make.jl
@@ -9,6 +9,7 @@ Pkg.instantiate()
 
 using Documenter
 using GNNLux
+using MLUtils
 using Lux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
 
diff --git a/GNNlib/docs/Project.toml b/GNNlib/docs/Project.toml
index 573ea3e88..f5aca8717 100644
--- a/GNNlib/docs/Project.toml
+++ b/GNNlib/docs/Project.toml
@@ -4,3 +4,4 @@ DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
 GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
 GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
diff --git a/GraphNeuralNetworks/docs/Project.toml b/GraphNeuralNetworks/docs/Project.toml
index a6f5fbee7..b4814d9e0 100644
--- a/GraphNeuralNetworks/docs/Project.toml
+++ b/GraphNeuralNetworks/docs/Project.toml
@@ -12,6 +12,7 @@ Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PlutoStaticHTML = "359b1769-a58e-495b-9770-312e911026ad"
 PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
diff --git a/GraphNeuralNetworks/docs/make.jl b/GraphNeuralNetworks/docs/make.jl
index 7090f479e..b82017c33 100644
--- a/GraphNeuralNetworks/docs/make.jl
+++ b/GraphNeuralNetworks/docs/make.jl
@@ -11,10 +11,12 @@ using Documenter
 using GraphNeuralNetworks
 using Flux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
+using MLUtils
 
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
 
 DocMeta.setdocmeta!(GNNGraphs, :DocTestSetup, :(using GNNGraphs, MLUtils); recursive = true)
+DocMeta.setdocmeta!(GNNlib, :DocTestSetup, :(using GNNlib); recursive = true)
 DocMeta.setdocmeta!(GraphNeuralNetworks, :DocTestSetup, :(using GraphNeuralNetworks); recursive = true)
 
 mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/mathtools"]),

From b296cfdc800f37494c89e34da63a8eb259252553 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 6 Feb 2025 08:02:33 +0100
Subject: [PATCH 8/8] cleanup

---
 GNNGraphs/docs/make.jl           | 1 -
 GNNLux/docs/make.jl              | 1 -
 GNNlib/docs/Project.toml         | 1 +
 GraphNeuralNetworks/docs/make.jl | 1 -
 4 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/GNNGraphs/docs/make.jl b/GNNGraphs/docs/make.jl
index ce02a7843..70f495946 100644
--- a/GNNGraphs/docs/make.jl
+++ b/GNNGraphs/docs/make.jl
@@ -6,7 +6,6 @@ Pkg.instantiate()
 using Documenter
 using DocumenterInterLinks
 using GNNGraphs
-using MLUtils # this is needed by setdocmeta!
 import Graphs
 using Graphs: induced_subgraph
 
diff --git a/GNNLux/docs/make.jl b/GNNLux/docs/make.jl
index 3c39c97f4..d79dd6171 100644
--- a/GNNLux/docs/make.jl
+++ b/GNNLux/docs/make.jl
@@ -9,7 +9,6 @@ Pkg.instantiate()
 
 using Documenter
 using GNNLux
-using MLUtils
 using Lux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
 
diff --git a/GNNlib/docs/Project.toml b/GNNlib/docs/Project.toml
index f5aca8717..3bd80885f 100644
--- a/GNNlib/docs/Project.toml
+++ b/GNNlib/docs/Project.toml
@@ -4,4 +4,5 @@ DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
 GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
 GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
diff --git a/GraphNeuralNetworks/docs/make.jl b/GraphNeuralNetworks/docs/make.jl
index b82017c33..d07c41d7a 100644
--- a/GraphNeuralNetworks/docs/make.jl
+++ b/GraphNeuralNetworks/docs/make.jl
@@ -11,7 +11,6 @@ using Documenter
 using GraphNeuralNetworks
 using Flux, GNNGraphs, GNNlib, Graphs
 using DocumenterInterLinks
-using MLUtils
 
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets