move to Flux 0.14 (#316)

* move to Flux 0.14

* cleanup

* bump CI

Author: CarloLucibello

.github/workflows/CI.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.7' # Replace this with the minimum Julia version that your package supports. E.g. if your package requires Julia 1.5 or higher, change this to '1.5'.
+          # - '1.9' # Replace this with the minimum Julia version that your package supports. E.g. if your package requires Julia 1.5 or higher, change this to '1.5'.
           - '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
           - 'nightly'
         os:

Project.toml

@@ -16,32 +16,32 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 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"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [compat]
 Adapt = "3"
-CUDA = "3.3, 4"
+CUDA = "4"
 ChainRulesCore = "1"
 DataStructures = "0.18"
-Flux = "0.13.9"
+Flux = "0.14"
 Functors = "0.4.1"
 Graphs = "1.4"
 KrylovKit = "0.6"
 MLDatasets = "0.7"
-MLUtils = "0.3, 0.4"
+MLUtils = "0.4"
 MacroTools = "0.5"
-NNlib = "0.8, 0.9"
-NNlibCUDA = "0.2"
+NNlib = "0.9"
 NearestNeighbors = "0.4"
 Reexport = "1"
-StatsBase = "0.33, 0.34"
-julia = "1.7"
+StatsBase = "0.34"
+cuDNN = "1"
+julia = "1.9"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/gnngraph.md

@@ -220,9 +220,9 @@ g′ = add_edges(g, [1, 2], [2, 3]) # add edges 1->2 and 2->3
 Move a `GNNGraph` to a CUDA device using `Flux.gpu` method. 
 
 ```julia
-using Flux: gpu
+using CUDA, Flux
 
-g_gpu = g |> gpu
+g_gpu = g |> Flux.gpu
 ```
 
 ## Integration with Graphs.jl

src/GNNGraphs/gnngraph.jl

@@ -66,7 +66,7 @@ functionality from that library.
 # Examples
 
 ```julia
-using Flux, GraphNeuralNetworks
+using Flux, GraphNeuralNetworks, CUDA
 
 # Construct from adjacency list representation
 data = [[2,3], [1,4,5], [1], [2,5], [2,4]]

src/GraphNeuralNetworks.jl

@@ -8,7 +8,7 @@ using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor, batch
 using MacroTools: @forward
 using MLUtils
-using NNlib, NNlibCUDA
+using NNlib
 using NNlib: scatter, gather
 using ChainRulesCore
 using Reexport

test/runtests.jl

@@ -1,9 +1,9 @@
+using CUDA
 using GraphNeuralNetworks
 using GraphNeuralNetworks.GNNGraphs: sort_edge_index
 using GraphNeuralNetworks.GNNGraphs: getn, getdata
 using Functors
 using Flux
-using CUDA
 using Flux: gpu, @functor
 using LinearAlgebra, Statistics, Random
 using NNlib