Merge pull request #250 from FluxML/fix

Resolve gradient bug for GatedGraphConv

Author: yuehhua

src/layers/conv.jl

@@ -324,13 +324,15 @@ function (ggc::GatedGraphConv)(fg::FeaturedGraph, H::AbstractMatrix{S}) where {T
     @assert (m <= ggc.out_ch) "number of input features must less or equals to output features."
     adj = adjacency_list(fg)
     if m < ggc.out_ch
-        Hpad = similar(H, S, ggc.out_ch - m, n)
-        H = vcat(H, fill!(Hpad, 0))
+        Hpad = Zygote.ignore() do
+            fill!(similar(H, S, ggc.out_ch - m, n), 0)
+        end
+        H = vcat(H, Hpad)
     end
     for i = 1:ggc.num_layers
         M = view(ggc.weight, :, :, i) * H
         _, M = propagate(ggc, adj, Fill(0.f0, 0, ne(fg)), M, +)
-        H, _ = ggc.gru(H, M)  # BUG: FluxML/Flux.jl#1381
+        H, _ = ggc.gru(H, M)
     end
     H
 end

test/layers/conv.jl

@@ -60,7 +60,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
             @test size(node_feature(fgt_)) == (out_channel, N)
 
             g = Zygote.gradient(x -> sum(node_feature(gc(x))), fg)[1]
-            @test size(g[].nf) == size(X)
+            @test size(g.nf) == size(X)
 
             g = Zygote.gradient(model -> sum(node_feature(model(fg))), gc)[1]
             @test size(g.weight) == size(gc.weight)
@@ -118,7 +118,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
             @test size(node_feature(fgt_)) == (out_channel, N)
 
             g = Zygote.gradient(x -> sum(node_feature(cc(x))), fg)[1]
-            @test size(g[].nf) == size(X)
+            @test size(g.nf) == size(X)
 
             g = Zygote.gradient(model -> sum(node_feature(model(fg))), cc)[1]
             @test size(g.weight) == size(cc.weight)
@@ -174,7 +174,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
             @test size(node_feature(fgt_)) == (out_channel, N)
 
             g = Zygote.gradient(x -> sum(node_feature(gc(x))), fg)[1]
-            @test size(g[].nf) == size(X)
+            @test size(g.nf) == size(X)
 
             g = Zygote.gradient(model -> sum(node_feature(model(fg))), gc)[1]
             @test size(g.weight1) == size(gc.weight1)
@@ -245,7 +245,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
                 @test size(node_feature(fgt_)) == (concat ? (out_channel*heads, N) : (out_channel, N))
 
                 g = Zygote.gradient(x -> sum(node_feature(gat(x))), fg_gat)[1]
-                @test size(g[].nf) == size(X)
+                @test size(g.nf) == size(X)
 
                 g = Zygote.gradient(model -> sum(node_feature(model(fg_gat))), gat)[1]
                 @test size(g.weight) == size(gat.weight)
@@ -299,7 +299,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
             @test size(node_feature(fgt_)) == (out_channel, N)
 
             g = Zygote.gradient(x -> sum(node_feature(ggc(x))), fg)[1]
-            @test size(g[].nf) == size(X)
+            @test size(g.nf) == size(X)
 
             g = Zygote.gradient(model -> sum(node_feature(model(fg))), ggc)[1]
             @test size(g.weight) == size(ggc.weight)
@@ -342,7 +342,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
             @test size(node_feature(fgt_)) == (out_channel, N)
 
             g = Zygote.gradient(x -> sum(node_feature(ec(x))), fg)[1]
-            @test size(g[].nf) == size(X)
+            @test size(g.nf) == size(X)
 
             g = Zygote.gradient(model -> sum(node_feature(model(fg))), ec)[1]
             @test size(g.nn.weight) == size(ec.nn.weight)
@@ -371,7 +371,7 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
 
             g = Zygote.gradient(x -> sum(node_feature(gc(x))), 
                                 FeaturedGraph(adj, nf=X))[1]
-            @test size(g.x.nf) == size(X)
+            @test size(g.nf) == size(X)
 
             g = Zygote.gradient(model -> sum(node_feature(model(FeaturedGraph(adj, nf=X)))), 
                                 gc)[1]

test/runtests.jl

@@ -1,12 +1,13 @@
 using GeometricFlux
 using GeometricFlux.Datasets
+using CUDA
 using Flux
 using Flux: @functor
 using FillArrays
 using GraphSignals
 using LightGraphs: SimpleGraph, SimpleDiGraph, add_edge!, nv, ne
 using LinearAlgebra
-using NNlib
+using NNlib, NNlibCUDA
 using SparseArrays: SparseMatrixCSC
 using Statistics: mean
 using Zygote
@@ -26,10 +27,7 @@ tests = [
     "models",
 ]
 
-if Flux.use_cuda[]
-    using CUDA
-    using Flux: gpu
-    using NNlibCUDA
+if CUDA.functional()
     append!(tests, cuda_tests)
 else
     @warn "CUDA unavailable, not testing GPU support"