fix some cuda problems

Author: CarloLucibello

src/GNNGraphs/convert.jl

@@ -29,9 +29,16 @@ function to_coo(A::SPARSE_T; dir=:out, num_nodes=nothing, weighted=true)
     return (s, t, v), num_nodes, num_edges
 end
 
-function to_coo(A::ADJMAT_T; dir=:out, num_nodes=nothing, weighted=true)
+function _findnz_idx(A)
     nz = findall(!=(0), A) # vec of cartesian indexes
     s, t = ntuple(i -> map(t->t[i], nz), 2)
+    return s, t, nz
+end
+
+@non_differentiable _findnz_idx(A)
+
+function to_coo(A::ADJMAT_T; dir=:out, num_nodes=nothing, weighted=true)
+    s, t, nz = _findnz_idx(A)
     v = A[nz] 
     if dir == :in
         s, t = t, s
@@ -175,7 +182,3 @@ function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing, weighted=
     end
     return A, num_nodes, num_edges
 end
-
-# @non_differentiable to_coo(x...)
-# @non_differentiable to_dense(x...)
-# @non_differentiable to_sparse(x...)

src/GNNGraphs/query.jl

@@ -127,26 +127,6 @@ end
 #     return [fneighs(g, i) for i in nodes]
 # end
 
-using ChainRulesCore: unthunk, NoTangent, ZeroTangent
-
-function ChainRulesCore.rrule(::typeof(findnz), A::AbstractSparseMatrix)
-    I, J, V = findnz(A)
-
-    function findnz_pullback(Δ)
-        Δ === NoTangent() && return (NoTangent(), Δ)
-        Δ === ZeroTangent() && return (NoTangent(), Δ)
-        
-        _, _, V̄ = unthunk(Δ)
-
-        V̄ === NoTangent() && return (NoTangent(), V̄)
-        V̄ === ZeroTangent() && return (NoTangent(), V̄)
-        
-        return NoTangent(), sparse(I, J, V̄)
-    end
-
-    return (I, J, V), findnz_pullback
-end
-
 adjacency_list(g::GNNGraph; dir=:out) = adjacency_list(g, 1:g.num_nodes; dir)
 
 

src/utils.jl

@@ -100,3 +100,13 @@ function broadcast_edges(g::GNNGraph, x)
     return gather(x, gi)
 end
 
+
+function ChainRulesCore.rrule(::typeof(Broadcast.broadcasted), T::Type{<:Number}, x::AbstractSparseArray)
+    proj = ProjectTo(x)
+
+    function broadcasted_cast_sparse(Δ)
+        return NoTangent(), NoTangent(), proj(Δ)         
+    end
+
+    return T.(x), broadcasted_cast_sparse
+end

test/GNNGraphs/convert.jl

@@ -0,0 +1,22 @@
+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 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
+        @test t isa CuVector
+        @test_broken s isa CuVector{Int32}
+        @test_broken t isa CuVector{Int32}
+        @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/query.jl

@@ -138,12 +138,8 @@
                 A = adjacency_matrix(g, weighted=true)
                 sum(A)
             end[1]
-            if GRAPH_T == :coo
-                # TODO use the @test option broken = (GRAPH_T != :coo) on julia >= 1.7
-                @test gw == [1,1,1]
-            else
-                @test_broken gw == [1,1,1]
-            end
+
+            @test gw == [1,1,1]
         end
     end
 end

test/runtests.jl

@@ -24,6 +24,7 @@ include("test_utils.jl")
 
 tests = [
     "GNNGraphs/gnngraph",
+    "GNNGraphs/convert",
     "GNNGraphs/transform",
     "GNNGraphs/operators",
     "GNNGraphs/generate",
@@ -32,20 +33,18 @@ tests = [
     "utils",
     "msgpass",
     "layers/basic",
-    "layers/conv",
-    "layers/pool",
-    "examples/node_classification_cora",
-    "deprecations",
+    # "layers/conv",
+    # "layers/pool",
+    # "examples/node_classification_cora",
+    # "deprecations",
 ]
 
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")
 
 @testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:dense, :coo, :sparse) 
     global GRAPH_T = graph_type
-    # global TEST_GPU = CUDA.functional() && (GRAPH_T != :sparse)
-    global TEST_GPU = false
-
-
+    global TEST_GPU = CUDA.functional() && (GRAPH_T != :sparse)
+    
     for t in tests
         startswith(t, "examples") && GRAPH_T == :dense && continue     # not testing :dense since causes OutOfMememory on github's CI
         include("$t.jl")