Global pooling of edge features in GlobalAttentionPool

In some situations it is nice to make global pooling of the edge features. In the implementation of GlobalAttentionPool
we have

function (l::GlobalAttentionPool)(g::GNNGraph, x::AbstractArray)
    α = softmax_nodes(g, l.fgate(x))
    feats = α .* l.ffeat(x)
    u = reduce_nodes(+, g, feats)
    return u   
end

As far as I can see we can easily get edge pooling by exchanging softmax_nodes with softmax_edges and reduce_nodes with reduce_edges.

Would it be possible to refactor the GlobalAttentionPool to accommodate this? If so I’ll be happy to contribute, but need a hint in the right direction. 😊

Start with the GlobalAttentionPool call shown in the issue and read how softmax_nodes, reduce_nodes, and their edge counterparts are implemented and used. Determine how the pool should support both node and edge features, then verify that the existing node behavior and the requested edge pooling both work correctly.