tests passing

Author: CarloLucibello

test/layers/conv.jl

@@ -27,7 +27,7 @@
     @testset "GCNConv" begin
         l = GCNConv(in_channel => out_channel)
         for g in test_graphs
-            gradtest(l, g)
+            gradtest(l, g, rtol=1e-5)
         end
 
         # l = GCNConv(in_channel => out_channel, relu, bias=false)
@@ -44,7 +44,7 @@
         @test size(l.bias) == (out_channel,)
         @test l.k == k
         for g in test_graphs
-            gradtest(l, g, rtol=1) #TODO broken
+            gradtest(l, g, rtol=1e-4, broken_grad_fields=[:weight])
         end
 
         @testset "bias=false" begin
@@ -56,7 +56,7 @@
     @testset "GraphConv" begin
         l = GraphConv(in_channel => out_channel)
         for g in test_graphs
-            gradtest(l, g, rtol=1e-3)
+            gradtest(l, g, rtol=1e-5)
         end
 
         # l = GraphConv(in_channel => out_channel, relu, bias=false)
@@ -72,18 +72,37 @@
 
     @testset "GATConv" begin
 
-        for heads = [1, 2], concat = [true, false]
-            l = GATConv(in_channel => out_channel; heads, concat)
-            for g in test_graphs
-                gradtest(l, g, atol=1e-3, rtol=1e-2) #TODO 
-            end
-    
-            # l = GraphConv(in_channel => out_channel, relu, bias=false)
-            # for g in test_graphs
-            #     gradtest(l, g)
-            # end
+        heads = 1
+        concat = true
+        l = GATConv(in_channel => out_channel; heads, concat)
+        for g in test_graphs
+            gradtest(l, g, rtol=1e-4)
+        end
+
+        heads = 2
+        concat = true
+        l = GATConv(in_channel => out_channel; heads, concat)
+        for g in test_graphs
+            gradtest(l, g, rtol=1e-4, 
+                broken_grad_fields = [:a])
         end
 
+        heads = 1
+        concat = false
+        l = GATConv(in_channel => out_channel; heads, concat)
+        for g in test_graphs
+            gradtest(l, g, rtol=1e-4, 
+                broken_grad_fields = [:a])
+        end
+
+        heads = 2
+        concat = false
+        l = GATConv(in_channel => out_channel; heads, concat)
+        gradtest(l, test_graphs[1], atol=1e-4, rtol=1e-4, 
+                    broken_grad_fields = [:a])
+        gradtest(l, test_graphs[2], atol=1e-4, rtol=1e-4)
+        
+
         @testset "bias=false" begin
             @test length(Flux.params(GATConv(2=>3))) == 3
             @test length(Flux.params(GATConv(2=>3, bias=false))) == 2
@@ -101,9 +120,9 @@
     end
 
     @testset "EdgeConv" begin
-        l = EdgeConv(Dense(2*in_channel, out_channel))
+        l = EdgeConv(Dense(2*in_channel, out_channel), aggr=+)
         for g in test_graphs
-            gradtest(l, g, atol=1e-5, rtol=1e-5) 
+            gradtest(l, g, rtol=1e-5)
         end
     end
 

test/runtests.jl

@@ -1,7 +1,7 @@
 using GraphNeuralNetworks
 using Flux
 using CUDA
-using Flux: gpu, @functor
+using Flux: gpu, @functor, f64, f32
 using LinearAlgebra, Statistics, Random
 using NNlib
 using LearnBase

test/test_utils.jl

@@ -2,7 +2,17 @@ using ChainRulesTestUtils, FiniteDifferences, Zygote
 
 const rule_config = Zygote.ZygoteRuleConfig()
 
-function gradtest(l, g::GNNGraph; atol=1e-9, rtol=1e-5)
+# Using this until https://github.com/JuliaDiff/FiniteDifferences.jl/issues/188
+# is fixed
+function FiniteDifferences.to_vec(x::Integer)
+    Integer_from_vec(v) = x
+    return Int[x], Integer_from_vec
+end
+
+function gradtest(l, g::GNNGraph; atol=1e-7, rtol=1e-5,
+                                 exclude_grad_fields=[],
+                                 broken_grad_fields=[]
+                                )
     # TODO these give errors, probably some bugs in ChainRulesTestUtils
     # test_rrule(rule_config, x -> l(g, x), x; rrule_f=rrule_via_ad, check_inferred=false)
     # test_rrule(rule_config, l -> l(g, x), l; rrule_f=rrule_via_ad, check_inferred=false)
@@ -27,27 +37,35 @@ function gradtest(l, g::GNNGraph; atol=1e-9, rtol=1e-5)
     # TEST LAYER GRADIENT - l(g, x) 
     l̄ = gradient(l -> sum(l(g, x)), l)[1]
     l̄_fd = FiniteDifferences.grad(fdm, l -> sum(l(g, x)), l)[1]
-    test_approx_structs(l, l̄, l̄_fd; atol, rtol)
-
+    test_approx_structs(l, l̄, l̄_fd; atol, rtol, broken_grad_fields, exclude_grad_fields)
     # TEST LAYER GRADIENT - l(g)
     l̄ = gradient(l -> sum(l(g).ndata.x), l)[1]
     l̄_fd = FiniteDifferences.grad(fdm, l -> sum(l(g).ndata.x), l)[1]
-    test_approx_structs(l, l̄, l̄_fd; atol, rtol)
+    test_approx_structs(l, l̄, l̄_fd; atol, rtol, broken_grad_fields, exclude_grad_fields)
 end
 
-function test_approx_structs(l, l̄, l̄_fd; atol=1e-9, rtol=1e-5)
+function test_approx_structs(l, l̄, l̄_fd; atol=1e-5, rtol=1e-5, 
+            broken_grad_fields=[],
+            exclude_grad_fields=[])
     for f in fieldnames(typeof(l))
+        f ∈ exclude_grad_fields && continue
         f̄, f̄_fd = getfield(l̄, f), getfield(l̄_fd, f)
         if isnothing(f̄)
-            @show f̄ f̄_fd
+            # @show f f̄_fd
             @test !(f̄_fd isa AbstractArray) || isapprox(f̄_fd, fill!(similar(f̄_fd), 0); atol=atol, rtol=rtol)
         elseif f̄ isa Union{AbstractArray, Number}
             @test eltype(f̄) == eltype(getfield(l, f))
-            @test f̄ ≈ f̄_fd   atol=atol rtol=rtol
+            if f ∈ broken_grad_fields
+                @test_broken f̄ ≈ f̄_fd   atol=atol rtol=rtol
+            else
+                # @show f getfield(l, f) f̄ f̄_fd broken_grad_fields
+                @test f̄ ≈ f̄_fd   atol=atol rtol=rtol
+            end
         else
-            test_approx_structs(getfield(l, f), f̄, f̄_fd)
+            test_approx_structs(getfield(l, f), f̄, f̄_fd; broken_grad_fields)
         end
     end
+    return true
 end