diff --git a/Project.toml b/Project.toml index b81560b53..d78ecb5cb 100644 --- a/Project.toml +++ b/Project.toml @@ -22,6 +22,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +SuiteSparseGraphBLAS = "c2e53296-7b14-11e9-1210-bddfa8111e1d" TestEnv = "1e6cf692-eddd-4d53-88a5-2d735e33781b" [compat] diff --git a/src/GNNGraphs/GNNGraphs.jl b/src/GNNGraphs/GNNGraphs.jl index 9177eb34f..d1cea7af0 100644 --- a/src/GNNGraphs/GNNGraphs.jl +++ b/src/GNNGraphs/GNNGraphs.jl @@ -14,6 +14,7 @@ using LearnBase: getobs import KrylovKit using ChainRulesCore using LinearAlgebra, Random +using SuiteSparseGraphBLAS: GBMatrix include("gnngraph.jl") export GNNGraph, diff --git a/src/GNNGraphs/convert.jl b/src/GNNGraphs/convert.jl index 0ce98ebcc..63d313d3f 100644 --- a/src/GNNGraphs/convert.jl +++ b/src/GNNGraphs/convert.jl @@ -124,7 +124,7 @@ function to_sparse(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothi A = T.(A) end if !(A isa AbstractSparseMatrix) - A = sparse(A) + A = GBMatrix(sparse(A)) end return A, num_nodes, num_edges end @@ -140,7 +140,7 @@ 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) - return A, num_nodes, num_edges + return GBMatrix(A), num_nodes, num_edges end diff --git a/test/runtests.jl b/test/runtests.jl index cfa8d9888..3153f6dab 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -37,7 +37,7 @@ tests = [ !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support") -@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo, :dense, :sparse) +@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:sparse, :coo, :dense,) global GRAPH_T = graph_type global TEST_GPU = CUDA.functional() && (GRAPH_T != :sparse)