Merge pull request #201 from mplemay/fix-flux-restructure

fix GNNChain restructure bug
JuliaGraphs · Jul 22, 2022 · 93b6fa2 · 93b6fa2
2 parents 401ae1a + c4d562c
commit 93b6fa2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -115,6 +115,7 @@ end
     Base.iterate, Base.lastindex, Base.keys
 
 Flux.functor(::Type{<:GNNChain}, c) = c.layers, ls -> GNNChain(ls...)
+Flux.functor(::Type{<:GNNChain}, c::Tuple) = c, ls -> GNNChain(ls...)
 
 # input from graph
 applylayer(l, g::GNNGraph) = GNNGraph(g, ndata=l(node_features(g)))

diff --git a/test/layers/basic.jl b/test/layers/basic.jl
@@ -64,5 +64,11 @@
         wg = WithGraph(model, g, traingraph=true)
         @test length(Flux.params(wg)) == length(Flux.params(model)) + length(Flux.params(g))
     end
+
+    @testset "Flux restructure" begin
+        chain = GNNChain(GraphConv(2=>2))
+        params, restructure = Flux.destructure(chain)
+        @test restructure(params) isa GNNChain
+    end
 end