separate GNNGraphs from GNNlib (#446)

* separate GNNGraphs from GNNlib

* complete factorization

* rebase

Author: CarloLucibello

GNNGraphs/Project.toml

@@ -0,0 +1,67 @@
+name = "GNNGraphs"
+uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c"
+authors = ["Carlo Lucibello and contributors"]
+version = "0.1.0"
+
+[deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LuxDeviceUtils = "34f89e08-e1d5-43b4-8944-0b49ac560553"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
+
+[extensions]
+GNNGraphsCUDAExt = "CUDA"
+GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
+
+[compat]
+Adapt = "4"
+CUDA = "5"
+ChainRulesCore = "1"
+Functors = "0.4.1"
+Graphs = "1.4"
+KrylovKit = "0.8"
+LinearAlgebra = "1"
+LuxDeviceUtils = "0.1.24"
+MLDatasets = "0.7"
+MLUtils = "0.4"
+NNlib = "0.9"
+NearestNeighbors = "0.4"
+Random = "1"
+SimpleWeightedGraphs = "1.4.0"
+SparseArrays = "1"
+Statistics = "1"
+StatsBase = "0.34"
+cuDNN = "1"
+julia = "1.9"
+
+[extras]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
+
+[targets]
+test = ["Test", "Adapt", "DataFrames", "InlineStrings", "SimpleWeightedGraphs", "Zygote", "FiniteDifferences", "ChainRulesTestUtils", "MLDatasets", "CUDA", "cuDNN"]

GNNGraphs/README.md

@@ -0,0 +1,7 @@
+# GNNGraphs.jl
+
+A package implementing graph types for graph deep learning. 
+
+This package is currently under development and may break frequentely. 
+It is not meant for final users but for GNN libraries developers.
+Final user should use GraphNeuralNetworks.jl instead.

GNNGraphs/ext/GNNGraphsCUDAExt/GNNGraphsCUDAExt.jl

@@ -0,0 +1,14 @@
+module GNNGraphsCUDAExt
+
+using CUDA
+using Random, Statistics, LinearAlgebra
+using GNNGraphs
+using GNNGraphs: COO_T, ADJMAT_T, SPARSE_T 
+
+const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}
+
+include("query.jl")
+include("transform.jl")
+include("utils.jl")
+
+end #module

GNNlib/ext/GNNlibCUDAExt/GNNGraphs/query.jl renamed to GNNGraphs/ext/GNNGraphsCUDAExt/query.jl

@@ -3,6 +3,9 @@ GNNGraphs.iscuarray(x::AnyCuArray) = true
 
 
 function sort_edge_index(u::AnyCuArray, v::AnyCuArray)
+    dev = get_device(u)
+    cdev = cpu_device()
+    u, v = u |> cdev, v |> cdev
     #TODO proper cuda friendly implementation
-    sort_edge_index(u |> Flux.cpu, v |> Flux.cpu) |> Flux.gpu
-end
+    sort_edge_index(u, v) |> dev
+end

GNNlib/ext/GNNlibCUDAExt/GNNGraphs/transform.jl renamed to GNNGraphs/ext/GNNGraphsCUDAExt/transform.jl

@@ -1,10 +1,10 @@
-module GNNlibSimpleWeightedGraphsExt
+module GNNGraphsSimpleWeightedGraphsExt
 
-using GNNlib
 using Graphs
+using GNNGraphs
 using SimpleWeightedGraphs
 
-function GNNlib.GNNGraph(g::T; kws...) where 
+function GNNGraphs.GNNGraph(g::T; kws...) where 
                  {T <: Union{SimpleWeightedGraph, SimpleWeightedDiGraph}}
    return GNNGraph(g.weights, kws...)
 end

GNNlib/ext/GNNlibCUDAExt/GNNGraphs/utils.jl renamed to GNNGraphs/ext/GNNGraphsCUDAExt/utils.jl

@@ -3,18 +3,18 @@ module GNNGraphs
 using SparseArrays
 using Functors: @functor
 import Graphs
-using Graphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree,
+using Graphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree, 
               has_self_loops, is_directed
-import MLUtils
-using MLUtils: getobs, numobs, ones_like, zeros_like, batch
 import NearestNeighbors
 import NNlib
 import StatsBase
 import KrylovKit
 using ChainRulesCore
 using LinearAlgebra, Random, Statistics
 import MLUtils
+using MLUtils: getobs, numobs, ones_like, zeros_like, chunk, batch
 import Functors
+using LuxDeviceUtils: get_device, cpu_device, LuxCPUDevice
 
 include("chainrules.jl") # hacks for differentiability
 
@@ -78,7 +78,9 @@ export add_nodes,
        to_unidirected,
        random_walk_pe,
        remove_nodes,
-# from Flux
+       ppr_diffusion,
+       drop_nodes,
+# from MLUtils
        batch,
        unbatch,
 # from SparseArrays
@@ -101,8 +103,7 @@ include("operators.jl")
 
 include("convert.jl")
 include("utils.jl")
-export sort_edge_index,
-       color_refinement
+export sort_edge_index, color_refinement
 
 include("gatherscatter.jl")
 # _gather, _scatter