update docstring

Author: CarloLucibello

docs/src/api/basic.md

@@ -14,7 +14,8 @@ Pages = ["basic.md"]
 
 ## Docs
 
-```@docs
-GNNLayer
-GNNChain
-```
+```@autodocs
+Modules = [GraphNeuralNetworks]
+Pages   = ["layers/basic.jl"]
+Private = false
+```

src/gnngraph.jl

@@ -305,8 +305,7 @@ function LightGraphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out)
         NNlib.scatter!(+, degs, src, s)
     end
     if dir ∈ [:in, :both]
-        # @show size(degs) src typeof(t)
-        NNlib.scatter!(+, degs, src, Int.(t))
+        NNlib.scatter!(+, degs, src, t)
     end
     return degs 
 end

src/layers/basic.jl

@@ -19,6 +19,9 @@ A type wrapping the `model` and tying it to the graph `g`.
 In the forward pass, can only take feature arrays as inputs,
 returning `model(g, x...; kws...)`.
 
+If `traingraph=false`, the graph's parameters, won't be collected
+when calling `Flux.params` on a `WithGraph` object.
+
 # Examples
 
 ```julia
@@ -41,11 +44,10 @@ struct WithGraph{M}
   traingraph::Bool
 end
 
-
 WithGraph(model, g::GNNGraph; traingraph=false) = WithGraph(model, g, traingraph)
 
 @functor WithGraph
-trainable(l::WithGraph) = l.traingraph ? (l.model, l.g) : (l.model,)
+Flux.trainable(l::WithGraph) = l.traingraph ? (l.model, l.g) : (l.model,)
 
 (l::WithGraph)(g::GNNGraph, x...; kws...) = l.model(g, x...; kws...)
 (l::WithGraph)(x...; kws...) = l.model(l.g, x...; kws...)

test/layers/basic.jl

@@ -46,17 +46,23 @@
     end
 
     @testset "WithGraph" begin
-        g = GNNGraph([1,2,3], [2,3,1])
         x = rand(Float32, 2, 3)
+        g = GNNGraph([1,2,3], [2,3,1], ndata=x)
         model = SAGEConv(2 => 3)
         wg = WithGraph(model, g)
         # No need to feed the graph to `wg`
         @test wg(x) == model(g, x)
-
+        @test Flux.params(wg) == Flux.params(model)
         g2 = GNNGraph([1,1,2,3], [2,4,1,1])
         x2 = rand(Float32, 2, 4)
         # WithGraph will ignore the internal graph if fed with a new one. 
         @test wg(g2, x2) == model(g2, x2)
+
+        wg = WithGraph(model, g, traingraph=false)
+        @test length(Flux.params(wg)) == length(Flux.params(model))
+
+        wg = WithGraph(model, g, traingraph=true)
+        @test length(Flux.params(wg)) == length(Flux.params(model)) + length(Flux.params(g))
     end
 end