add softmax and readout

Author: CarloLucibello

docs/make.jl

@@ -18,7 +18,7 @@ makedocs(;
                 "Convolutional Layers" => "api/conv.md",
                 "Pooling Layers" => "api/pool.md",
                 "Message Passing" => "api/messagepassing.md",
-                "NNlib" => "api/nnlib.md",
+                "Utils" => "api/utils.md",
                ],
               "Developer Notes" => "dev.md",
             ],

docs/src/api/nnlib.md

@@ -0,0 +1,35 @@
+```@meta
+CurrentModule = GraphNeuralNetworks
+```
+
+# Utility Functions
+
+## Index
+
+```@index
+Order = [:type, :function]
+Pages   = ["utils.md"]
+```
+
+## Docs
+
+
+### Readout Functions 
+
+```@docs
+GraphNeuralNetworks.readout_nodes
+GraphNeuralNetworks.readout_edges
+GraphNeuralNetworks.softmax_nodes
+GraphNeuralNetworks.softmax_edges
+```
+
+### NNlib
+
+Primitive functions implemented in NNlib.jl.
+
+```@docs
+NNlib.gather!
+NNlib.gather
+NNlib.scatter!
+NNlib.scatter
+```

docs/src/api/utils.md

@@ -30,6 +30,10 @@ export
     # from SparseArrays
     sprand, sparse, blockdiag,
 
+    # utils
+    readout_nodes, readout_edges, 
+    softmax_nodes, softmax_edges,
+    
     # msgpass
     apply_edges, propagate,
     copyxj,

src/GraphNeuralNetworks.jl

@@ -129,4 +129,68 @@ function NNlib.scatter!(op, dst::AnyCuArray, src::Number, idx::AnyCuArray)
     blocks = cld(max_idx, threads)
     kernel(args...; threads=threads, blocks=blocks)
     return dst
+end
+
+"""
+    readout_nodes(g, x, aggr)
+
+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)
+    indexes = graph_indicator(g)
+    return NNlib.scatter(aggr, x, indexes)
+end
+
+"""
+    readout_edges(g, e, aggr)
+
+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)
+    s, t = edge_index(g)
+    indexes = graph_indicator(g)[s]
+    return NNlib.scatter(aggr, e, indexes)
+end
+
+"""
+    softmax_nodes(g, x, aggr)
+
+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
+    gi = graph_indicator(g)
+    den = den[:, gi]
+    return num ./ reshape(den, size(num))
+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
+    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))
+end
+
+function graph_indicator(g)
+    if isnothing(g.graph_indicator)
+        return ones_like(edge_index(g)[1], Int, g.num_nodes)
+    else 
+        return g.graph_indicator
+    end
 end

src/utils.jl

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

test/runtests.jl

@@ -0,0 +1,30 @@
+@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]) 
+    x = g.ndata.x
+    e = g.edata.e
+
+    @testset "readout_nodes" begin
+        r = readout_nodes(g, x, mean)
+        @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)
+        @test size(r) == (De, g.num_graphs)
+        @test r[:,2] ≈ mean(getgraph(g, 2).edata.e, dims=2)
+    end
+
+    @testset "softmax_nodes" begin
+        r = softmax_nodes(g, x)
+        @test size(r) == size(x)
+        @test r[:,1:10] ≈ softmax(getgraph(g, 1).ndata.x, dims=2)
+    end
+
+    @testset "softmax_edges" begin
+        r = softmax_edges(g, e)
+        @test size(r) == size(e)
+        @test r[:,1:60] ≈ softmax(getgraph(g, 1).edata.e, dims=2)
+    end
+end