Replace bypass_graph as Bypass layer

Author: yuehhua

src/GeometricFlux.jl

@@ -54,11 +54,8 @@ export
     InnerProductDecoder,
     VariationalEncoder,
 
-    # layer/selector
-    bypass_graph,
-
-    # utils
-    generate_cluster,
+    # layer/misc
+    Bypass,
 
     #node2vec
     node2vec

src/layers/misc.jl

@@ -1,13 +1,32 @@
 """
-    bypass_graph(nf_func, ef_func, gf_func)
-
-Bypassing graph in FeaturedGraph and let other layer process (node, edge and global)features only.
+Bypassing graph in FeaturedGraph and let other layer process (node, edge and global) features only.
 """
-function bypass_graph(nf_func=identity, ef_func=identity, gf_func=identity)
-    return function (fg::FeaturedGraph)
-        FeaturedGraph(fg,
-                      nf=nf_func(node_feature(fg)),
-                      ef=ef_func(edge_feature(fg)),
-                      gf=gf_func(global_feature(fg)))
-    end
+struct Bypass{N,E,G}
+    node_layer::N
+    edge_layer::E
+    global_layer::G
+end
+
+@functor Bypass
+
+Bypass(; node_layer=identity, edge_layer=identity, global_layer=identity) =
+    Bypass(node_layer, edge_layer, global_layer)
+
+function (l::Bypass)(fg::FeaturedGraph)
+    nf = l.node_layer(node_feature(fg))
+    ef = l.edge_layer(edge_feature(fg))
+    gf = l.global_layer(global_feature(fg))
+    return FeaturedGraph(fg, nf=nf, ef=ef, gf=gf)
+end
+
+function (l::Bypass)(fsg::FeaturedSubgraph)
+    nf = l.node_layer(node_feature(fsg))
+    ef = l.edge_layer(edge_feature(fsg))
+    gf = l.global_layer(global_feature(fsg))
+    fg = parent(fsg)
+    vidx = fsg.nodes
+    nf = NNlib.scatter(+, nf, vidx; init=0, dstsize=(size(nf,1), nv(fg)))
+    ef = NNlib.scatter(+, ef, edges(fsg); init=0, dstsize=(size(ef,1), ne(fg)))
+    fg = FeaturedGraph(fg, nf=nf, ef=ef, gf=gf)
+    return subgraph(fg, vidx)
 end

test/layers/misc.jl

@@ -1,24 +1,26 @@
 @testset "misc" begin
-    @testset "bypass_graph" begin
+    @testset "Bypass" begin
+        T = Float32
         N = 4
         E = 5
-        adj = [0 1 1 1;
-            1 0 1 0;
-            1 1 0 1;
-            1 0 1 0]
+        adj = T[0 1 1 1;
+                1 0 1 0;
+                1 1 0 1;
+                1 0 1 0]
 
         nf = rand(3, N)
         ef = rand(5, E)
         gf = rand(7)
+        nodes = [1,2,4]
 
         fg = FeaturedGraph(adj, nf=nf, ef=ef, gf=gf)
-        layer = bypass_graph(x -> x .+ 1.,
-                                x -> x .+ 2.,
-                                x -> x .+ 3.)
-        fg_ = layer(fg)
-        @test GraphSignals.adjacency_matrix(fg_) == adj
-        @test node_feature(fg_) == nf .+ 1.
-        @test edge_feature(fg_) == ef .+ 2.
-        @test global_feature(fg_) == gf .+ 3.
+        fsg = subgraph(fg, nodes)
+
+        layer = Bypass(node_layer=Dropout(0.5),
+                       global_layer=x -> x .+ 3.)
+        fsg_ = layer(fsg)
+        @test node_feature(fsg_) == view(nf, :, nodes)
+        @test edge_feature(fsg_) == view(ef, :, edges(fsg))
+        @test global_feature(fsg_) == gf .+ 3.
     end
 end