improve degree and more test fixes (#268)

* fix gpu test + test rearrangment

edge_weight in rand_graph

drop julia 1.6

improve degree

* fix the rules

* fix

* rrule for adjacency matrix

* cleanup

* cleanup

* ops

* all tests passing also on gpu

Author: CarloLucibello

.github/workflows/CI.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6' # 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.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' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
           - 'nightly'
         os:

Project.toml

@@ -41,7 +41,7 @@ NNlibCUDA = "0.2"
 NearestNeighbors = "0.4"
 Reexport = "1"
 StatsBase = "0.33"
-julia = "1.6"
+julia = "1.7"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/GNNGraphs/generate.jl

@@ -1,13 +1,15 @@
 """
-    rand_graph(n, m; bidirected=true, seed=-1, kws...)
+    rand_graph(n, m; bidirected=true, seed=-1, edge_weight = nothing, kws...)
 
-Generate a random (Erdós-Renyi) `GNNGraph` with `n` nodes
-and `m` edges.
+Generate a random (Erdós-Renyi) `GNNGraph` with `n` nodes and `m` edges.
 
 If `bidirected=true` the reverse edge of each edge will be present.
 If `bidirected=false` instead, `m` unrelated edges are generated.
 In any case, the output graph will contain no self-loops or multi-edges.
 
+A vector can be passed  as `edge_weight`. Its length has to be equal to `m`
+in the directed case, and `m÷2` in the bidirected one.
+
 Use a `seed > 0` for reproducibility.
 
 Additional keyword arguments will be passed to the [`GNNGraph`](@ref) constructor.
@@ -37,12 +39,12 @@ julia> edge_index(g)
 
 ```
 """
-function rand_graph(n::Integer, m::Integer; bidirected = true, seed = -1, kws...)
+function rand_graph(n::Integer, m::Integer; bidirected = true, seed = -1, edge_weight = nothing, kws...)
     if bidirected
         @assert iseven(m) "Need even number of edges for bidirected graphs, given m=$m."
     end
     m2 = bidirected ? m ÷ 2 : m
-    return GNNGraph(Graphs.erdos_renyi(n, m2; is_directed = !bidirected, seed); kws...)
+    return GNNGraph(Graphs.erdos_renyi(n, m2; is_directed = !bidirected, seed); edge_weight, kws...)
 end
 
 """
@@ -68,7 +70,7 @@ GNNHeteroGraph:
   num_edges: ((:user, :rate, :movie) => 30,)
 ```
 """
-rand_heteropraph
+function rand_heteropraph end
 
 # for generic iterators of pairs
 rand_heterograph(n, m; kws...) = rand_heterograph(Dict(n), Dict(m); kws...)

src/GNNGraphs/gnngraph.jl

@@ -176,15 +176,20 @@ GNNGraph((s, t)::NTuple{2}; kws...) = GNNGraph((s, t, nothing); kws...)
 
 # GNNGraph(g::AbstractGraph; kws...) = GNNGraph(adjacency_matrix(g, dir=:out); kws...)
 
-function GNNGraph(g::AbstractGraph; kws...)
+function GNNGraph(g::AbstractGraph; edge_weight = nothing, kws...)
     s = Graphs.src.(Graphs.edges(g))
     t = Graphs.dst.(Graphs.edges(g))
+    w = edge_weight
     if !Graphs.is_directed(g)
         # add reverse edges since GNNGraph is directed
         s, t = [s; t], [t; s]
+        if !isnothing(w)
+            @assert length(w) == Graphs.ne(g) "edge_weight must have length equal to the number of undirected edges"
+            w = [w; w]
+        end
     end
     num_nodes::Int = Graphs.nv(g)
-    GNNGraph((s, t); num_nodes = num_nodes, kws...)
+    GNNGraph((s, t, w); num_nodes = num_nodes, kws...)
 end
 
 function GNNGraph(g::GNNGraph; ndata = g.ndata, edata = g.edata, gdata = g.gdata,

src/GNNGraphs/query.jl

@@ -144,7 +144,8 @@ If `weighted=true`, the `A` will contain the edge weights if any, otherwise the
 function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType = eltype(g); dir = :out,
                                  weighted = true)
     if g.graph[1] isa CuVector
-        # TODO revisit after https://github.com/JuliaGPU/CUDA.jl/pull/1152
+        # Revisit after 
+        # https://github.com/JuliaGPU/CUDA.jl/issues/1113
         A, n, m = to_dense(g.graph, T; num_nodes = g.num_nodes, weighted)
     else
         A, n, m = to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted)
@@ -164,63 +165,101 @@ function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType = eltype(g
     return dir == :out ? A : A'
 end
 
-function _get_edge_weight(g, edge_weight)
-    if edge_weight === true || edge_weight === nothing
-        ew = get_edge_weight(g)
-    elseif edge_weight === false
-        ew = nothing
-    elseif edge_weight isa AbstractVector
-        ew = edge_weight
+function ChainRulesCore.rrule(::typeof(adjacency_matrix), g::G, T::DataType; 
+            dir = :out, weighted = true) where {G <: GNNGraph{<:ADJMAT_T}}
+    A = adjacency_matrix(g, T; dir, weighted)
+    if !weighted
+        function adjacency_matrix_pullback_noweight(Δ)
+            return (NoTangent(), ZeroTangent(), NoTangent())  
+        end
+        return A, adjacency_matrix_pullback_noweight
     else
-        error("Invalid edge_weight argument.")
+        function adjacency_matrix_pullback_weighted(Δ)
+            dg = Tangent{G}(; graph = Δ .* binarize(A))
+            return (NoTangent(), dg, NoTangent())  
+        end
+        return A, adjacency_matrix_pullback_weighted
+    end
+end
+
+function ChainRulesCore.rrule(::typeof(adjacency_matrix), g::G, T::DataType; 
+            dir = :out, weighted = true) where {G <: GNNGraph{<:COO_T}}
+    A = adjacency_matrix(g, T; dir, weighted)
+    w = get_edge_weight(g)
+    if !weighted || w === nothing
+        function adjacency_matrix_pullback_noweight(Δ)
+            return (NoTangent(), ZeroTangent(), NoTangent())  
+        end
+        return A, adjacency_matrix_pullback_noweight
+    else
+        function adjacency_matrix_pullback_weighted(Δ)
+            s, t = edge_index(g)
+            dg = Tangent{G}(; graph = (NoTangent(), NoTangent(), NNlib.gather(Δ, s, t)))
+            return (NoTangent(), dg, NoTangent())  
+        end
+        return A, adjacency_matrix_pullback_weighted
+    end
+end
+
+function _get_edge_weight(g, edge_weight::Bool)
+    if edge_weight === true
+        return get_edge_weight(g)
+    elseif edge_weight === false
+        return nothing
     end
-    return ew
 end
 
+_get_edge_weight(g, edge_weight::AbstractVector) = edge_weight
+
 """
     degree(g::GNNGraph, T=nothing; dir=:out, edge_weight=true)
 
 Return a vector containing the degrees of the nodes in `g`.
 
+The gradient is propagated through this function only if `edge_weight` is `true`
+or a vector.
+
 # Arguments
+
 - `g`: A graph.
 - `T`: Element type of the returned vector. If `nothing`, is
        chosen based on the graph type and will be an integer
-       if `edge_weight=false`.
+       if `edge_weight=false`. Default `nothing`.
 - `dir`: For `dir=:out` the degree of a node is counted based on the outgoing edges.
          For `dir=:in`, the ingoing edges are used. If `dir=:both` we have the sum of the two.
 - `edge_weight`: If `true` and the graph contains weighted edges, the degree will 
                 be weighted. Set to `false` instead to just count the number of
-                outgoing/ingoing edges.
-                In alternative, you can also pass a vector of weights to be used
+                outgoing/ingoing edges. 
+                Finally, you can also pass a vector of weights to be used
                 instead of the graph's own weights.
+                Default `true`.
+
 """
 function Graphs.degree(g::GNNGraph{<:COO_T}, T::TT = nothing; dir = :out,
                        edge_weight = true) where {
                                                   TT <: Union{Nothing, Type{<:Number}}}
     s, t = edge_index(g)
 
-    edge_weight = _get_edge_weight(g, edge_weight)
-    edge_weight = edge_weight === nothing ? ones_like(s) : edge_weight
-
-    T = isnothing(T) ? eltype(edge_weight) : T
-    degs = fill!(similar(s, T, g.num_nodes), 0)
-
-    if dir ∈ [:out, :both]
-        degs = degs .+ NNlib.scatter(+, edge_weight, s, dstsize = (g.num_nodes,))
-    end
-    if dir ∈ [:in, :both]
-        degs = degs .+ NNlib.scatter(+, edge_weight, t, dstsize = (g.num_nodes,))
-    end
-    return degs
+    ew = _get_edge_weight(g, edge_weight)
+    
+    T = if isnothing(T)
+            if !isnothing(ew)
+                eltype(ew)
+            else
+                eltype(s)
+            end
+        else 
+            T
+        end
+    return _degree((s, t), T, dir, ew, g.num_nodes)
 end
 
 # TODO:: Make efficient
 Graphs.degree(g::GNNGraph, i::Union{Int, AbstractVector}; dir = :out) = degree(g; dir)[i]
 
 function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T::TT = nothing; dir = :out,
-                       edge_weight = true) where {TT}
-    TT <: Union{Nothing, Type{<:Number}}
+                       edge_weight = true) where {TT<:Union{Nothing, Type{<:Number}}}
+    
     # edge_weight=true or edge_weight=nothing act the same here
     @assert !(edge_weight isa AbstractArray) "passing the edge weights is not support by adjacency matrix representations"
     @assert dir ∈ (:in, :out, :both)
@@ -234,6 +273,26 @@ function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T::TT = nothing; dir = :out,
         end
     end
     A = adjacency_matrix(g)
+    return _degree(A, T, dir, edge_weight, g.num_nodes)
+end
+
+function _degree((s, t)::Tuple, T::Type, dir::Symbol, edge_weight::Nothing, num_nodes::Int)
+    _degree((s, t), T, dir, ones_like(s, T), num_nodes)
+end
+
+function _degree((s, t)::Tuple, T::Type, dir::Symbol, edge_weight::AbstractVector, num_nodes::Int)
+    degs = fill!(similar(s, T, num_nodes), 0)
+
+    if dir ∈ [:out, :both]
+        degs = degs .+ NNlib.scatter(+, edge_weight, s, dstsize = (num_nodes,))
+    end
+    if dir ∈ [:in, :both]
+        degs = degs .+ NNlib.scatter(+, edge_weight, t, dstsize = (num_nodes,))
+    end
+    return degs
+end
+
+function _degree(A::AbstractMatrix, T::Type, dir::Symbol, edge_weight::Bool, num_nodes::Int)
     if edge_weight === false
         A = binarize(A)
     end
@@ -243,6 +302,40 @@ function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T::TT = nothing; dir = :out,
            vec(sum(A, dims = 1)) .+ vec(sum(A, dims = 2))
 end
 
+function ChainRulesCore.rrule(::typeof(_degree), graph, T, dir, edge_weight::Nothing, num_nodes)
+    degs = _degree(graph, T, dir, edge_weight, num_nodes)
+    function _degree_pullback(Δ)
+        return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())
+    end
+    return degs, _degree_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(_degree), A::ADJMAT_T, T, dir, edge_weight::Bool, num_nodes)
+    degs = _degree(A, T, dir, edge_weight, num_nodes)
+    if edge_weight === false
+        function _degree_pullback_noweights(Δ)
+            return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())
+        end
+        return degs, _degree_pullback_noweights
+    else
+        function _degree_pullback_weights(Δ)
+            # We propagate the gradient only to the non-zero elements
+            # of the adjacency matrix.
+            bA = binarize(A)
+            if dir == :in
+                dA = bA .* Δ'
+            elseif dir == :out
+                dA = Δ .* bA
+            else # dir == :both
+                dA = Δ .* bA + Δ' .* bA
+            end
+            return (NoTangent(), dA, NoTangent(), NoTangent(), NoTangent(), NoTangent())
+        end
+        return degs, _degree_pullback_weights
+    end
+end
+
+
 """
     has_isolated_nodes(g::GNNGraph; dir=:out)
 

src/layers/conv.jl

@@ -26,7 +26,17 @@ and optionally an edge weight vector.
 - `init`: Weights' initializer. Default `glorot_uniform`.
 - `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
 - `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available).
-                     If `add_self_loops=true` the new weights will be set to 1. Default `false`.
+                     If `add_self_loops=true` the new weights will be set to 1. 
+                     This option is ignored if the `edge_weight` is explicitly provided in the forward pass.
+                     Default `false`.
+
+# Forward
+
+    (::GCNConv)(g::GNNGraph, x::AbstractMatrix, edge_weight = nothing) -> AbstractMatrix
+
+Takes as input a graph `g`,ca node feature matrix `x` of size `[in, num_nodes]`,
+and optionally an edge weight vector. Returns a node feature matrix of size 
+`[out, num_nodes]`.
 
 # Examples
 
@@ -107,7 +117,11 @@ function (l::GCNConv)(g::GNNGraph,
         # multiply before convolution if it is more convenient, otherwise multiply after
         x = l.weight * x
     end
-    d = degree(g, T; dir = :in, edge_weight)
+    if edge_weight !== nothing
+        d = degree(g, T; dir = :in, edge_weight)
+    else
+        d = degree(g, T; dir = :in, edge_weight = l.use_edge_weight)
+    end
     c = 1 ./ sqrt.(d)
     x = x .* c'
     if edge_weight !== nothing
@@ -1288,7 +1302,11 @@ function (l::SGConv)(g::GNNGraph, x::AbstractMatrix{T},
     if Dout < Din
         x = l.weight * x
     end
-    d = degree(g, T; dir = :in, edge_weight)
+    if edge_weight !== nothing
+        d = degree(g, T; dir = :in, edge_weight)
+    else
+        d = degree(g, T; dir = :in, edge_weight=l.use_edge_weight)
+    end
     c = 1 ./ sqrt.(d)
     for iter in 1:(l.k)
         x = x .* c'

test/GNNGraphs/convert.jl

@@ -10,10 +10,8 @@ if TEST_GPU
         @test Array(y) ≈ [2, 2, 2, 2, 2, 2]
 
         s, t = get_st(A)
-        @test s isa CuVector
-        @test t isa CuVector
-        @test_broken s isa CuVector{Int32}
-        @test_broken t isa CuVector{Int32}
+        @test s isa CuVector{<:Integer}
+        @test t isa CuVector{<:Integer}
         @test Array(s) == [2, 3, 1, 3, 1, 2]
         @test Array(t) == [1, 1, 2, 2, 3, 3]
 

test/GNNGraphs/generate.jl

@@ -22,6 +22,14 @@
 
     g2 = rand_graph(n, m, bidirected = false, seed = 17, graph_type = GRAPH_T)
     @test edge_index(g2) == edge_index(g)
+
+    ew = rand(m2)
+    g = rand_graph(n, m, bidirected = true, seed = 17, graph_type = GRAPH_T, edge_weight = ew)
+    @test get_edge_weight(g) == [ew; ew] broken=(GRAPH_T != :coo)
+    
+    ew = rand(m)
+    g = rand_graph(n, m, bidirected = false, seed = 17, graph_type = GRAPH_T, edge_weight = ew)
+    @test get_edge_weight(g) == ew broken=(GRAPH_T != :coo)
 end
 
 @testset "knn_graph" begin

test/GNNGraphs/gnnheterograph.jl

@@ -90,4 +90,3 @@ end
 #     @test sprint(show, MIME("text/plain"), hg3; context=:compact => false) =="GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)\n ndata:\n\t:A  =>  (x = 2×10 Matrix{Float64}, y = 3×10 Matrix{Float64})\n\t:B  =>  x = 10×20 Matrix{Float64}"
 #     @test sprint(show, MIME("text/plain"), hg2; context=:compact => false) != sprint(show, MIME("text/plain"), hg3; context=:compact => false)
 # end
-