cleanup

Author: CarloLucibello

Project.toml

@@ -5,7 +5,6 @@ version = "0.2.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -14,7 +13,6 @@ KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LearnBase = "7f8f8fb0-2700-5f03-b4bd-41f8cfc144b6"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
@@ -28,7 +26,7 @@ Adapt = "3"
 CUDA = "3.3"
 ChainRulesCore = "1"
 DataStructures = "0.18"
-Flux = "0.12"
+Flux = "0.12.7"
 KrylovKit = "0.5"
 LearnBase = "0.4, 0.5"
 LightGraphs = "1.3"

src/deprecations.jl

@@ -16,14 +16,14 @@ update_edge(l, e, m) = e
 function propagate(l::GNNLayer, g::GNNGraph, aggr, x, e=nothing)
     @warn """
           Passing a GNNLayer to propagate is deprecated, 
-          you should pass directly the message function.
+          you should pass the message function directly.
           The new signature is `propagate(f, g, aggr; [xi, xj, e])`.
 
-          Also the functions `compute_message`, `update_node`,
-          and `update_edge` have been deprecated. Please
+          The functions `compute_message`, `update_node`,
+          and `update_edge` have been deprecated as well. Please
           refer to the documentation.
           """
-    m = apply_edge((a...) -> compute_message(l, a...), g, x, x, e)
+    m = apply_edges((a...) -> compute_message(l, a...), g, x, x, e)
     m̄ = aggregate_neighbors(g, aggr, m)
     x = update_node(l, x, m̄)
     e = update_edge(l, e, m)

src/utils.jl

@@ -74,14 +74,6 @@ ones_like(x::CUMAT_T, T=eltype(x), sz=size(x)) = CUDA.ones(T, sz)
 
 ofeltype(x, y) = convert(float(eltype(x)), y)
 
-# TODO move to flux. fix for https://github.com/FluxML/Flux.jl/issues/1720
-Flux._cpu_array(x::AbstractSparseArray) = Flux.adapt(SparseMatrixCSC, x)
-
-# TODO. FIX THIS HACK. CUDA.jl support to sparse matrices is very bad, convert to dense
-# Revisit after https://github.com/JuliaGPU/CUDA.jl/pull/1152
-Flux._gpu_array(x::AbstractSparseArray) = CuMatrix(x)
-
-
 # Considers the src a zero dimensional object.
 # Useful for implementing `StatsBase.counts`, `degree`, etc...
 # function NNlib.scatter!(op, dst::AbstractArray, src::Number, idx::AbstractArray)

test/deprecations.jl

@@ -0,0 +1,34 @@
+@testset "deprecations" begin
+    @testset "propagate" begin
+        struct GCN{A<:AbstractMatrix, B, F} <: GNNLayer
+            weight::A
+            bias::B
+            σ::F
+        end
+
+        Flux.@functor GCN # allow collecting params, gpu movement, etc...
+
+        function GCN(ch::Pair{Int,Int}, σ=identity)
+            in, out = ch
+            W = Flux.glorot_uniform(out, in)
+            b = zeros(Float32, out)
+            GCN(W, b, σ)
+        end
+
+        GraphNeuralNetworks.compute_message(l::GCN, xi, xj, e) = xj 
+
+        function (l::GCN)(g::GNNGraph, x::AbstractMatrix{T}) where T
+            x, _ = propagate(l, g, +, x) 
+            return l.σ.(l.weight * x .+ l.bias)
+        end
+
+        function new_forward(l, g, x)
+            x = propagate(copyxj, g, +, xj=x) 
+            return l.σ.(l.weight * x .+ l.bias)
+        end
+
+        g = GNNGraph(random_regular_graph(10, 4), ndata=randn(3, 10))
+        l = GCN(3 => 5, tanh)
+        @test l(g, g.ndata.x) ≈ new_forward(l, g, g.ndata.x)
+    end
+end

test/runtests.jl

@@ -24,6 +24,7 @@ tests = [
     "layers/conv",
     "layers/pool",
     "examples/node_classification_cora",
+    "deprecations",
 ]
 
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")