add nodes/edges broadcast

Author: CarloLucibello

docs/src/api/utils.md

@@ -14,13 +14,15 @@ Pages   = ["utils.md"]
 ## Docs
 
 
-### Readout Functions 
+### Graph-wise operations 
 
 ```@docs
-GraphNeuralNetworks.readout_nodes
-GraphNeuralNetworks.readout_edges
+GraphNeuralNetworks.reduce_nodes
+GraphNeuralNetworks.reduce_edges
 GraphNeuralNetworks.softmax_nodes
 GraphNeuralNetworks.softmax_edges
+GraphNeuralNetworks.broadcast_nodes
+GraphNeuralNetworks.broadcast_edges
 ```
 
 ### NNlib

src/GraphNeuralNetworks.jl

@@ -12,6 +12,7 @@ using MacroTools: @forward
 import LearnBase
 using LearnBase: getobs
 using NNlib, NNlibCUDA
+using NNlib: scatter, gather
 using ChainRulesCore
 import LightGraphs
 using LightGraphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree
@@ -31,8 +32,9 @@ export
     sprand, sparse, blockdiag,
 
     # utils
-    readout_nodes, readout_edges, 
+    reduce_nodes, reduce_edges, 
     softmax_nodes, softmax_edges,
+    broadcast_nodes, broadcast_edges,
 
     # msgpass
     apply_edges, propagate,

src/layers/pool.jl

@@ -34,7 +34,7 @@ struct GlobalPool{F} <: GNNLayer
 end
 
 function (l::GlobalPool)(g::GNNGraph, x::AbstractArray)
-    return readout_nodes(g, x, l.aggr)
+    return reduce_nodes(g, x, l.aggr)
 end
 
 """

src/utils.jl

@@ -132,65 +132,97 @@ function NNlib.scatter!(op, dst::AnyCuArray, src::Number, idx::AnyCuArray)
 end
 
 """
-    readout_nodes(g, x, aggr)
+    reduce_nodes(aggr, g, x)
 
 For a batched graph `g`, return the graph-wise aggregation of the node
 features `x`. The aggregation operator `aggr` can be `+`, `mean`, `max`, or `min`.
 The returned array will have last dimension `g.num_graphs`.
 """
-function readout_nodes(g, x, aggr)
+function reduce_nodes(aggr, g::GNNGraph, x)
+    @assert size(x)[end] == g.num_nodes
     indexes = graph_indicator(g)
     return NNlib.scatter(aggr, x, indexes)
 end
 
 """
-    readout_edges(g, e, aggr)
+    reduce_edges(aggr, g, e)
 
 For a batched graph `g`, return the graph-wise aggregation of the edge
 features `e`. The aggregation operator `aggr` can be `+`, `mean`, `max`, or `min`.
 The returned array will have last dimension `g.num_graphs`.
 """
-function readout_edges(g, e, aggr)
+function reduce_edges(aggr, g::GNNGraph, e)
+    @assert size(e)[end] == g.num_edges
     s, t = edge_index(g)
     indexes = graph_indicator(g)[s]
     return NNlib.scatter(aggr, e, indexes)
 end
 
 """
-    softmax_nodes(g, x, aggr)
+    softmax_nodes(g, x)
 
 Graph-wise softmax of the node features `x`.
 """
-function softmax_nodes(g, x)
-    max_ = maximum(x; dims = ndims(x)) # TODO use graph-wise maximum
-    num = exp.(x .- max_)
-    den = readout_nodes(g, num, +)
-    den = Flux.flatten(den) # reshape to matrix for convenience
+function softmax_nodes(g::GNNGraph, x)
+    @assert size(x)[end] == g.num_nodes
     gi = graph_indicator(g)
-    den = den[:, gi]
-    return num ./ reshape(den, size(num))
+    max_ = gather(scatter(max, x, gi), gi)
+    num = exp.(x .- max_)
+    den = reduce_nodes(+, g, num)
+    den = gather(den, gi)
+    return num ./ den
 end
 
 """
     softmax_edges(g, e)
 
 Graph-wise softmax of the edge features `e`.
 """
-function softmax_edges(g, e)
-    max_ = maximum(e; dims = ndims(e)) # TODO use graph-wise maximum
+function softmax_edges(g::GNNGraph, e)
+    @assert size(e)[end] == g.num_edges
+    gi = graph_indicator(g, edges=true)
+    max_ = gather(scatter(max, e, gi), gi)
     num = exp.(e .- max_)
-    den = readout_edges(g, num, +)
-    den = Flux.flatten(den) # reshape to matrix for convenience
-    s, t = edge_index(g)
-    gi = graph_indicator(g)[s]
-    den = den[:, gi]
-    return num ./ reshape(den, size(num))
+    den = reduce_edges(+, g, num)
+    den = gather(den, gi)
+    return num ./ den
+end
+
+"""
+    broadcast_nodes(g, x)
+
+Graph-wise broadcast array `x` of size `(*, g.num_graphs)` 
+to size `(*, g.num_nodes)`.
+"""
+function broadcast_nodes(g::GNNGraph, x)
+    @assert size(x)[end] == g.num_graphs
+    gi = graph_indicator(g)
+    return gather(x, gi)
 end
 
-function graph_indicator(g)
+"""
+    broadcast_edges(g, x)
+
+Graph-wise broadcast array `x` of size `(*, g.num_graphs)` 
+to size `(*, g.num_edges)`.
+"""
+function broadcast_edges(g::GNNGraph, x)
+    @assert size(x)[end] == g.num_graphs
+    gi = graph_indicator(g, edges=true)
+    return gather(x, gi)
+end
+
+
+function graph_indicator(g; edges=false)
     if isnothing(g.graph_indicator)
-        return ones_like(edge_index(g)[1], Int, g.num_nodes)
+        gi = ones_like(edge_index(g)[1], Int, g.num_nodes)
     else 
-        return g.graph_indicator
+        gi = g.graph_indicator
+    end
+    if edges
+        s, t = edge_index(g)
+        return gi[s]
+    else
+        return gi
     end
-end
+end

test.jl

@@ -0,0 +1,70 @@
+using Flux, Random
+
+
+function get_grad1(m, data)
+    gradient(Flux.params(m)) do
+        loss(m, data)
+    end
+end
+
+function get_grad2(m, data)
+    gradient(Flux.params(m)) do
+        ps = Flux.params(m) # just creating params without using them
+        loss(m, data)
+    end
+end
+
+function get_grad3(m, data)
+    gradient(Flux.params(m)) do
+        ps = Flux.params(m)
+        loss(m, data) + sum(sum(p) for p in ps)
+    end
+end
+
+function get_grad4(m, data)
+    ps = Flux.params(m)
+    gradient(Flux.params(m)) do
+        loss(m, data) + sum(sum(p) for p in ps)
+    end
+end
+
+function get_grad5(m, data)
+    gradient(Flux.params(m)) do
+        sum(Flux.params(m)[1]) + sum(Flux.params(m)[2])
+    end
+end
+
+function get_grad6(m, data)
+    ps = Flux.params(m)
+    gradient(Flux.params(m)) do
+        sum(ps[1]) + sum(ps[2])
+    end
+end
+
+function get_grad7(m, data)
+    ps = Flux.params(m)
+    gradient(Flux.params(m)) do
+        sum(m.weight) + sum(m.bias)
+    end
+end
+
+Random.seed!(17)
+m = Dense(3, 2);
+data = rand(Float32, 3, 5)
+loss(m, x) = sum(m(x).^2)
+
+g1 = get_grad1(m, data) 
+g2 = get_grad2(m, data)
+g3 = get_grad3(m, data) 
+g4 = get_grad4(m, data) 
+g5 = get_grad5(m, data) 
+g6 = get_grad6(m, data) 
+g7 = get_grad7(m, data) 
+
+@show g1[m.weight] #            correct
+@show g2[m.weight] # == g1 .+ 1  wrong, should be == g1
+@show g3[m.weight] # == g1 .+ 2  wrong, should be == g1 .+ 1
+@show g4[m.weight] # == g1 .+ 1  correct
+@show g5[m.weight] # .== 2       wrong, should be .== 1
+@show g6[m.weight] # == nothing  wrong,  should be .== 1
+@show g7[m.weight] # == 1        correct

test/runtests.jl

@@ -19,7 +19,7 @@ include("test_utils.jl")
 
 tests = [
     "gnngraph",
-    "utils"
+    "utils",
     "msgpass",
     "layers/basic",
     "layers/conv",

test/utils.jl

@@ -1,17 +1,20 @@
 @testset "Utils" begin
     De, Dx = 3, 2
-    g = Flux.batch([GNNGraph(erdos_renyi(10, 30), ndata=rand(Dx, 10), edata=rand(De, 30)) for i=1:5]) 
+    g = Flux.batch([GNNGraph(erdos_renyi(10, 30), 
+                             ndata=rand(Dx, 10), 
+                             edata=rand(De, 30),
+                             graph_type=GRAPH_T) for i=1:5]) 
     x = g.ndata.x
     e = g.edata.e
 
-    @testset "readout_nodes" begin
-        r = readout_nodes(g, x, mean)
+    @testset "reduce_nodes" begin
+        r = reduce_nodes(mean, g, x)
         @test size(r) == (Dx, g.num_graphs)
         @test r[:,2] ≈ mean(getgraph(g, 2).ndata.x, dims=2)
     end
 
-    @testset "readout_edges" begin
-        r = readout_edges(g, e, mean)
+    @testset "reduce_edges" begin
+        r = reduce_edges(mean, g, e)
         @test size(r) == (De, g.num_graphs)
         @test r[:,2] ≈ mean(getgraph(g, 2).edata.e, dims=2)
     end
@@ -27,4 +30,23 @@
         @test size(r) == size(e)
         @test r[:,1:60] ≈ softmax(getgraph(g, 1).edata.e, dims=2)
     end
-end
+
+
+    @testset "broadcast_nodes" begin
+        z = rand(4, g.num_graphs)
+        r = broadcast_nodes(g, z)
+        @test size(r) == (4, g.num_nodes)
+        @test r[:,1] ≈ z[:,1]
+        @test r[:,10] ≈ z[:,1]
+        @test r[:,11] ≈ z[:,2]
+    end
+
+    @testset "broadcast_edges" begin
+        z = rand(4, g.num_graphs)
+        r = broadcast_edges(g, z)
+        @test size(r) == (4, g.num_edges)
+        @test r[:,1] ≈ z[:,1]
+        @test r[:,60] ≈ z[:,1]
+        @test r[:,61] ≈ z[:,2]
+    end
+end