fixing overflow in GAT layers (#244)

* fixing overflow in GAT layers

* one more small fix

* an updated version of fix of GAT network

Author: pevnak

src/layers/conv.jl

@@ -343,8 +343,11 @@ function (l::GATConv)(g::GNNGraph, x::AbstractMatrix, e::Union{Nothing,AbstractM
     Wx = l.dense_x(x)
     Wx = reshape(Wx, chout, heads, :)                   # chout × nheads × nnodes
 
-    m = propagate(message, g, +, l; xi=Wx, xj=Wx, e)                 ## chout × nheads × nnodes
-    x = m.β ./ m.α
+    # a hand-writtent message passing
+    m = apply_edges((xi, xj, e) -> message(l, xi, xj, e), g, Wx, Wx, e)
+    α = softmax_edge_neighbors(g, m.logα)
+    β = α .* m.Wxj
+    x = aggregate_neighbors(g, +, β)
 
     if !l.concat
         x = mean(x, dims=2)
@@ -367,8 +370,8 @@ function message(l::GATConv, Wxi, Wxj, e)
         Wxx = vcat(Wxi, Wxj, We)
     end
     aWW = sum(l.a .* Wxx, dims=1)   # 1 × nheads × nedges
-    α = exp.(leakyrelu.(aWW, l.negative_slope))       
-    return (α = α, β = α .* Wxj)
+    logα = leakyrelu.(aWW, l.negative_slope)
+    return(logα = logα, Wxj = Wxj)
 end
 
 function Base.show(io::IO, l::GATConv)

src/utils.jl

@@ -54,7 +54,7 @@ function softmax_edges(g::GNNGraph, e)
     num = exp.(e .- max_)
     den = reduce_edges(+, g, num)
     den = gather(den, gi)
-    return num ./ den
+    return num ./ (den .+ eps(eltype(e)))
 end
 
 @doc raw"""