diff --git a/src/layers/conv.jl b/src/layers/conv.jl index 092a1241b..dbb094da7 100644 --- a/src/layers/conv.jl +++ b/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 diff --git a/test/layers/conv.jl b/test/layers/conv.jl index 0e0f165b2..8482858de 100644 --- a/test/layers/conv.jl +++ b/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] diff --git a/test/runtests.jl b/test/runtests.jl index 46f335f15..e3e26b598 100644 --- a/test/runtests.jl +++ b/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"