temporarily reintegrate GNNGraphs tests (#449)

* add back GNNGraph tests

* reintegrate extensions

* don't use DeviceUtils

* import Flux

Author: CarloLucibello

ext/GraphNeuralNetworksCUDAExt/GNNGraphs/query.jl

@@ -0,0 +1,2 @@
+
+GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1))

ext/GraphNeuralNetworksCUDAExt/GNNGraphs/transform.jl

@@ -0,0 +1,2 @@
+
+GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz)

ext/GraphNeuralNetworksCUDAExt/GNNGraphs/utils.jl

@@ -0,0 +1,8 @@
+
+GNNGraphs.iscuarray(x::AnyCuArray) = true
+
+
+function sort_edge_index(u::AnyCuArray, v::AnyCuArray)
+    #TODO proper cuda friendly implementation
+    sort_edge_index(u |> Flux.cpu, v |> Flux.cpu) |> Flux.gpu
+end

ext/GraphNeuralNetworksCUDAExt/GraphNeuralNetworksCUDAExt.jl

@@ -9,6 +9,9 @@ import GraphNeuralNetworks: propagate
 
 const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}
 
+include("GNNGraphs/query.jl")
+include("GNNGraphs/transform.jl")
+include("GNNGraphs/utils.jl")
 include("msgpass.jl")
 
 end #module

ext/GraphNeuralNetworksSimpleWeightedGraphsExt/GraphNeuralNetworksSimpleWeightedGraphsExt.jl

@@ -0,0 +1,12 @@
+module GraphNeuralNetworksSimpleWeightedGraphsExt
+
+using GraphNeuralNetworks
+using Graphs
+using SimpleWeightedGraphs
+
+function GraphNeuralNetworks.GNNGraph(g::T; kws...) where 
+                 {T <: Union{SimpleWeightedGraph, SimpleWeightedDiGraph}}
+   return GNNGraph(g.weights, kws...)
+end
+
+end #module

src/GNNGraphs/GNNGraphs.jl

@@ -7,6 +7,8 @@ using Graphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree
               has_self_loops, is_directed
 import NearestNeighbors
 import NNlib
+import Flux
+using Flux: batch
 import StatsBase
 import KrylovKit
 using ChainRulesCore

src/GNNGraphs/transform.jl

@@ -1126,10 +1126,13 @@ function negative_sample(g::GNNGraph;
 
     s, t = edge_index(g)
     n = g.num_nodes
-    device = get_device(s)
-    cdevice = cpu_device()
-    # Convert to gpu since set operations and sampling are not supported by CUDA.jl
-    s, t = cdevice(s), cdevice(t)
+    if iscuarray(s)
+        # Convert to gpu since set operations and sampling are not supported by CUDA.jl
+        device = Flux.gpu
+        s, t = Flux.cpu(s), Flux.cpu(t)
+    else
+        device = Flux.cpu
+    end
     idx_pos, maxid = edge_encoding(s, t, n)
     if bidirected
         num_neg_edges = num_neg_edges ÷ 2
@@ -1156,6 +1159,7 @@ function negative_sample(g::GNNGraph;
     return GNNGraph(s_neg, t_neg, num_nodes = n) |> device
 end
 
+
 """
     rand_edge_split(g::GNNGraph, frac; bidirected=is_bidirected(g)) -> g1, g2
 

test/GNNGraphs/chainrules.jl

@@ -0,0 +1,24 @@
+@testset "dict constructor" begin
+    grad = gradient(1.) do x
+        d = Dict([:x => x, :y => 5]...)    
+        return sum(d[:x].^2)
+    end[1]
+
+    @test grad == 2
+
+    ## BROKEN Constructors
+    # grad = gradient(1.) do x
+    #     d = Dict([(:x => x), (:y => 5)])    
+    #     return sum(d[:x].^2)
+    # end[1]
+
+    # @test grad == 2
+
+
+    # grad = gradient(1.) do x
+    #     d = Dict([(:x => x), (:y => 5)])    
+    #     return sum(d[:x].^2)
+    # end[1]
+
+    # @test grad == 2
+end

test/GNNGraphs/convert.jl

@@ -0,0 +1,20 @@
+if TEST_GPU
+    @testset "to_coo(dense) on gpu" begin
+        get_st(A) = GNNGraphs.to_coo(A)[1][1:2]
+        get_val(A) = GNNGraphs.to_coo(A)[1][3]
+
+        A = cu([0 2 2; 2.0 0 2; 2 2 0])
+
+        y = get_val(A)
+        @test y isa CuVector{Float32}
+        @test Array(y) ≈ [2, 2, 2, 2, 2, 2]
+
+        s, t = get_st(A)
+        @test s isa CuVector{<:Integer}
+        @test t isa CuVector{<:Integer}
+        @test Array(s) == [2, 3, 1, 3, 1, 2]
+        @test Array(t) == [1, 1, 2, 2, 3, 3]
+
+        @test gradient(A -> sum(get_val(A)), A)[1] isa CuMatrix{Float32}
+    end
+end

test/GNNGraphs/datastore.jl

@@ -0,0 +1,101 @@
+
+@testset "constructor" begin
+    @test_throws AssertionError DataStore(10, (:x => rand(10), :y => rand(2, 4)))
+
+    @testset "keyword args" begin
+        ds = DataStore(10, x = rand(10), y = rand(2, 10))
+        @test size(ds.x) == (10,)
+        @test size(ds.y) == (2, 10)
+
+        ds = DataStore(x = rand(10), y = rand(2, 10))
+        @test size(ds.x) == (10,)
+        @test size(ds.y) == (2, 10)
+    end
+end
+
+@testset "getproperty / setproperty!" begin
+    x = rand(10)
+    ds = DataStore(10, (:x => x, :y => rand(2, 10)))
+    @test ds.x == ds[:x] == x
+    @test_throws DimensionMismatch ds.z=rand(12)
+    ds.z = [1:10;]
+    @test ds.z == [1:10;]
+    vec = [DataStore(10, (:x => x,)), DataStore(10, (:x => x, :y => rand(2, 10)))]
+    @test vec.x == [x, x]
+    @test_throws KeyError vec.z 
+    @test vec._n == [10, 10]
+    @test vec._data == [Dict(:x => x), Dict(:x => x, :y => vec[2].y)]
+end
+
+@testset "setindex!" begin
+    ds = DataStore(10)
+    x = rand(10)
+    @test (ds[:x] = x) == x # Tests setindex!
+    @test ds.x == ds[:x] == x
+end
+
+@testset "map" begin
+    ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
+    ds2 = map(x -> x .+ 1, ds)
+    @test ds2.x == ds.x .+ 1
+    @test ds2.y == ds.y .+ 1
+
+    @test_throws AssertionError ds2=map(x -> [x; x], ds)
+end
+
+@testset "getdata / getn" begin
+    ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
+    @test getdata(ds) == getfield(ds, :_data)
+    @test_throws KeyError ds.data
+    @test getn(ds) == getfield(ds, :_n)
+    @test_throws KeyError ds.n
+end
+
+@testset "cat empty" begin
+    ds1 = DataStore(2, (:x => rand(2)))
+    ds2 = DataStore(1, (:x => rand(1)))
+    dsempty = DataStore(0, (:x => rand(0)))
+
+    ds = GNNGraphs.cat_features(ds1, ds2)
+    @test getn(ds) == 3
+    ds = GNNGraphs.cat_features(ds1, dsempty)
+    @test getn(ds) == 2
+
+    # issue #280
+    g = GNNGraph([1], [2])
+    h = add_edges(g, Int[], Int[])  # adds no edges
+    @test getn(g.edata) == 1
+    @test getn(h.edata) == 1
+end
+
+
+@testset "gradient" begin
+    ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
+
+    f1(ds) = sum(ds.x)
+    grad = gradient(f1, ds)[1]
+    @test grad._data[:x] ≈ ngradient(f1, ds)[1][:x]
+
+    g = rand_graph(5, 2)
+    x = rand(2, 5)
+    grad = gradient(x -> sum(exp, GNNGraph(g, ndata = x).ndata.x), x)[1]
+    @test grad == exp.(x)
+end
+
+@testset "functor" begin
+    ds = DataStore(10, (:x => zeros(10), :y => ones(2, 10)))
+    p, re = Functors.functor(ds)
+    @test p[1] === getn(ds)
+    @test p[2] === getdata(ds)
+    @test ds == re(p)
+
+    ds2 = Functors.fmap(ds) do x
+        if x isa AbstractArray
+            x .+ 1
+        else
+            x
+        end
+    end
+    @test ds isa DataStore
+    @test ds2.x == ds.x .+ 1
+end