Add GATConv support for HeteroGraphConv (#400)

* gat hetero support

* Update src/layers/conv.jl

Co-authored-by: Carlo Lucibello <[email protected]>

* changes made

* fix

* Update src/layers/conv.jl

Co-authored-by: Carlo Lucibello <[email protected]>

---------

Co-authored-by: Carlo Lucibello <[email protected]>

Author: rbSparky

src/layers/conv.jl

@@ -409,25 +409,36 @@ end
 
 (l::GATConv)(g::GNNGraph) = GNNGraph(g, ndata = l(g, node_features(g), edge_features(g)))
 
-function (l::GATConv)(g::GNNGraph, x::AbstractMatrix,
+function (l::GATConv)(g::AbstractGNNGraph, x,
                       e::Union{Nothing, AbstractMatrix} = nothing)
     check_num_nodes(g, x)
     @assert !((e === nothing) && (l.dense_e !== nothing)) "Input edge features required for this layer"
     @assert !((e !== nothing) && (l.dense_e === nothing)) "Input edge features were not specified in the layer constructor"
 
+    xj, xi = expand_srcdst(g, x)
+
     if l.add_self_loops
         @assert e===nothing "Using edge features and setting add_self_loops=true at the same time is not yet supported."
-        g = add_self_loops(g)
+       if g isa GNNHeteroGraph
+           g = add_self_loops(g,  g.etypes[1])
+       else
+          g = add_self_loops(g)
+       end
     end
 
     _, chout = l.channel
     heads = l.heads
 
-    Wx = l.dense_x(x)
-    Wx = reshape(Wx, chout, heads, :)                   # chout × nheads × nnodes
+    Wxi = Wxj = l.dense_x(xj)
+    Wxi = Wxj = reshape(Wxj, chout, heads, :)                   
+
+    if xi !== xj
+        Wxi = l.dense_x(xi)
+        Wxi = reshape(Wxi, chout, heads, :)                   
+    end
 
     # a hand-written message passing
-    m = apply_edges((xi, xj, e) -> message(l, xi, xj, e), g, Wx, Wx, e)
+    m = apply_edges((xi, xj, e) -> message(l, xi, xj, e), g, Wxi, Wxj, e)
     α = softmax_edge_neighbors(g, m.logα)
     β = α .* m.Wxj
     x = aggregate_neighbors(g, +, β)

src/utils.jl

@@ -81,8 +81,14 @@ Softmax over each node's neighborhood of the edge features `e`.
                     {\sum_{j'\in N(i)} e^{\mathbf{e}_{j'\to i}}}.
 ```
 """
-function softmax_edge_neighbors(g::GNNGraph, e)
-    @assert size(e)[end] == g.num_edges
+function softmax_edge_neighbors(g::AbstractGNNGraph, e)
+    if g isa GNNHeteroGraph
+        for (key, value) in g.num_edges
+            @assert size(e)[end] == value
+        end
+    else
+        @assert size(e)[end] == g.num_edges
+    end
     s, t = edge_index(g)
     max_ = gather(scatter(max, e, t), t)
     num = exp.(e .- max_)

test/layers/heteroconv.jl

@@ -125,6 +125,14 @@
         @test size(y.A) == (2, 2) && size(y.B) == (2, 3)
     end
 
+    @testset "GATConv" begin
+        x = (A = rand(Float32, 4, 2), B = rand(Float32, 4, 3))
+        layers = HeteroGraphConv((:A, :to, :B) => GATConv(4 => 2),
+                                 (:B, :to, :A) => GATConv(4 => 2));
+        y = layers(hg, x); 
+        @test size(y.A) == (2, 2) && size(y.B) == (2, 3)
+    end
+
     @testset "GINConv" begin
         x = (A = rand(4, 2), B = rand(4, 3))
         layers = HeteroGraphConv((:A, :to, :B) => GINConv(Dense(4, 2), 0.4),