clenaup

Author: CarloLucibello

Project.toml

@@ -4,6 +4,7 @@ authors = ["Carlo Lucibello and contributors"]
 version = "0.1.1"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"

src/layers/conv.jl

@@ -340,7 +340,7 @@ update_node(l::GatedGraphConv, m, x) = m
 # remove after https://github.com/JuliaDiff/ChainRules.jl/pull/521
 @non_differentiable fill!(x...)
 
-function (l::GatedGraphConv)(g::GNNGraph, H::AbstractMatrix{S}) where {T<:AbstractVector,S<:Real}
+function (l::GatedGraphConv)(g::GNNGraph, H::AbstractMatrix{S}) where {S<:Real}
     check_num_nodes(g, H)
     m, n = size(H)
     @assert (m <= l.out_ch) "number of input features must less or equals to output features."

test/layers/conv.jl

@@ -46,14 +46,11 @@
         @test size(l.bias) == (out_channel,)
         @test l.k == k
         for g in test_graphs
-            if g === g_single_vertex && GRAPH_T == :dense
-                @test_broken test_layer(l, g, rtol=1e-5, broken_grad_fields=[:weight], test_gpu=false)
-            else
-                test_layer(l, g, rtol=1e-5, broken_grad_fields=[:weight], test_gpu=false)
-                if TEST_GPU
-                    @test_broken test_layer(l, g, rtol=1e-5, broken_grad_fields=[:weight], test_gpu=true)
-                end
-            end            
+            g = add_self_loops(g)
+            test_layer(l, g, rtol=1e-5, test_gpu=false, outsize=(out_channel, g.num_nodes))
+            if TEST_GPU
+                @test_broken test_layer(l, g, rtol=1e-5, test_gpu=true, outsize=(out_channel, g.num_nodes))
+            end              
         end
 
         @testset "bias=false" begin
@@ -81,10 +78,10 @@
 
     @testset "GATConv" begin
 
-        for heads in (1, 3), concat in (true, false)
+        for heads in (1, 2), concat in (true, false)
             l = GATConv(in_channel => out_channel; heads, concat)
             for g in test_graphs
-                test_layer(l, g, rtol=1e-4, 
+                test_layer(l, g, rtol=1e-4,
                     outsize=(concat ? heads*out_channel : out_channel, g.num_nodes))
             end
         end