add WithGraph

Author: CarloLucibello

src/GraphNeuralNetworks.jl

@@ -43,6 +43,7 @@ export
     # layers/basic
     GNNLayer,
     GNNChain,
+    WithGraph,
 
     # layers/conv
     CGConv,

src/gnngraph.jl

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

src/layers/basic.jl

@@ -11,6 +11,46 @@ abstract type GNNLayer end
 # To be specialized by layers also needing edge features as input (e.g. NNConv). 
 (l::GNNLayer)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g)))
 
+
+"""
+  WithGraph(model, g::GNNGraph; traingraph=false) 
+
+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...)`.
+
+# Examples
+
+```julia
+g = GNNGraph([1,2,3], [2,3,1])
+x = rand(Float32, 2, 3)
+model = SAGEConv(2 => 3)
+wg = WithGraph(model, g)
+# No need to feed the graph to `wg`
+@assert wg(x) == model(g, x)
+
+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. 
+@assert wg(g2, x2) == model(g2, x2)
+```
+"""
+struct WithGraph{M}
+  model::M
+  g::GNNGraph
+  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,)
+
+(l::WithGraph)(g::GNNGraph, x...; kws...) = l.model(g, x...; kws...)
+(l::WithGraph)(x...; kws...) = l.model(l.g, x...; kws...)
+
+
 """
     GNNChain(layers...)
     GNNChain(name = layer, ...)

test/layers/basic.jl

@@ -44,5 +44,19 @@
             @test chain(g).ndata.x == y
         end
     end
+
+    @testset "WithGraph" begin
+        g = GNNGraph([1,2,3], [2,3,1])
+        x = rand(Float32, 2, 3)
+        model = SAGEConv(2 => 3)
+        wg = WithGraph(model, g)
+        # No need to feed the graph to `wg`
+        @test wg(x) == model(g, x)
+
+        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)
+    end
 end