use GBMatrix

Author: CarloLucibello

Project.toml

@@ -17,6 +17,7 @@ NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+SuiteSparseGraphBLAS = "c2e53296-7b14-11e9-1210-bddfa8111e1d"
 
 [compat]
 CUDA = "3.3"

src/GraphNeuralNetworks.jl

@@ -14,6 +14,7 @@ using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
 using MacroTools: @forward
 using NNlib, NNlibCUDA
 using ChainRulesCore
+using SuiteSparseGraphBLAS: GBMatrix
 import LightGraphs
 using LightGraphs: AbstractGraph, outneighbors, inneighbors, is_directed, ne, nv, 
                   adjacency_matrix, degree

src/graph_conversions.jl

@@ -111,7 +111,7 @@ function to_sparse(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothi
     if T != eltype(A)
         A = T.(A)
     end
-    return sparse(A), num_nodes, num_edges
+    return GBMatrix(sparse(A)), num_nodes, num_edges
 end
 
 function to_sparse(adj_list::ADJLIST_T, T::DataType=Int; dir=:out, num_nodes=nothing)
@@ -125,7 +125,8 @@ function to_sparse(coo::COO_T, T::DataType=Int; dir=:out, num_nodes=nothing)
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     A = sparse(s, t, eweight, num_nodes, num_nodes)
     num_edges = length(s)
-    A, num_nodes, num_edges
+    
+    GBMatrix(A), num_nodes, num_edges
 end
 
 @non_differentiable to_coo(x...)

test/runtests.jl

@@ -23,7 +23,9 @@ 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, :sparse, :dense)
+# @testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo, :sparse, :dense)
+@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:sparse,)
+
     global GRAPH_T = graph_type
     for t in tests
         include("$t.jl")