fix gpu tests

Author: CarloLucibello

src/msgpass.jl

@@ -146,10 +146,10 @@ copyxj(xi, xj, e) = xj
 # ximulxj(xi, xj, e) = xi .* xj
 # xiaddxj(xi, xj, e) = xi .+ xj
 
-function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e)
-    A = adjacency_matrix(g)
-    return xj * A
-end
+# function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e)
+#     A = adjacency_matrix(g)
+#     return xj * A
+# end
 
 # function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(mean), xi, xj::AbstractMatrix, e)
 #     A = adjacency_matrix(g)

test/layers/basic.jl

@@ -15,7 +15,7 @@
 
         testmode!(gnn)
 
-        test_layer(gnn, g, rtol=1e-5)
+        test_layer(gnn, g, rtol=1e-5, exclude_grad_fields=[:μ, :σ²])
 
 
         @testset "Parallel" begin
@@ -29,7 +29,7 @@
 
             testmode!(gnn)
 
-            test_layer(gnn, g, rtol=1e-5)
+            test_layer(gnn, g, rtol=1e-5, exclude_grad_fields=[:μ, :σ²])
         end
     end
 end

test/runtests.jl

@@ -27,7 +27,7 @@ tests = [
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")
 
 # Testing all graph types. :sparse is a bit broken at the moment
-@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo,)
+@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo,:sparse,:dense)
 
     global GRAPH_T = graph_type
     global TEST_GPU = CUDA.functional() && GRAPH_T != :sparse

test/test_utils.jl

@@ -99,7 +99,6 @@ function test_layer(l, g::GNNGraph; atol = 1e-7, rtol = 1e-5,
 
     # TEST LAYER GRADIENT - l(g, x) 
     l̄ = gradient(l -> loss(l, g, x), l)[1]
-    l̄ = l̄ isa Base.RefValue ? l̄[] : l̄           # Zygote wraps gradient of mutables in RefValue 
     l̄_fd = FiniteDifferences.grad(fdm, l64 -> loss(l64, g64, x64), l64)[1]
     test_approx_structs(l, l̄, l̄_fd; atol, rtol, broken_grad_fields, exclude_grad_fields, verbose)
 
@@ -110,7 +109,6 @@ function test_layer(l, g::GNNGraph; atol = 1e-7, rtol = 1e-5,
 
     # TEST LAYER GRADIENT - l(g)
     l̄ = gradient(l -> loss(l, g), l)[1]
-    l̄ = l̄ isa Base.RefValue ? l̄[] : l̄           # Zygote wraps gradient of mutables in RefValue 
     l̄_fd = FiniteDifferences.grad(fdm, l64 -> loss(l64, g64), l64)[1]
     test_approx_structs(l, l̄, l̄_fd; atol, rtol, broken_grad_fields, exclude_grad_fields, verbose)
 
@@ -122,6 +120,9 @@ function test_approx_structs(l, l̄, l̄2; atol=1e-5, rtol=1e-5,
             exclude_grad_fields=[],
             verbose=false)
 
+    l̄ = l̄ isa Base.RefValue ? l̄[] : l̄           # Zygote wraps gradient of mutables in RefValue 
+    l̄2 = l̄2 isa Base.RefValue ? l̄2[] : l̄2           # Zygote wraps gradient of mutables in RefValue 
+
     for f in fieldnames(typeof(l))
         f ∈ exclude_grad_fields && continue
         f̄, f̄2 = getfield(l̄, f), getfield(l̄2, f)
@@ -147,7 +148,6 @@ function test_approx_structs(l, l̄, l̄2; atol=1e-5, rtol=1e-5,
             end
         else
             verbose && println("C")
-            f̄ = f̄ isa Base.RefValue ? f̄[] : f̄          # Zygote wraps gradient of mutables in RefValue 
             test_approx_structs(x, f̄, f̄2; exclude_grad_fields, broken_grad_fields, verbose)
         end
     end