Add GATv2Conv support to HeteroGraphConv (#407)

* hetero support

* Update conv.jl

Author: rbSparky

src/layers/conv.jl

@@ -579,21 +579,23 @@ end
 
 (l::GATv2Conv)(g::GNNGraph) = GNNGraph(g, ndata = l(g, node_features(g), edge_features(g)))
 
-function (l::GATv2Conv)(g::GNNGraph, x::AbstractMatrix,
+function (l::GATv2Conv)(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)
     end
     _, out = l.channel
     heads = l.heads
 
-    Wxi = reshape(l.dense_i(x), out, heads, :)                                  # out × heads × nnodes
-    Wxj = reshape(l.dense_j(x), out, heads, :)                                  # out × heads × nnodes
+    Wxi = reshape(l.dense_i(xi), out, heads, :)                                  # out × heads × nnodes
+    Wxj = reshape(l.dense_j(xj), out, heads, :)                                  # out × heads × nnodes
 
     m = apply_edges((xi, xj, e) -> message(l, xi, xj, e), g, Wxi, Wxj, e)
     α = softmax_edge_neighbors(g, m.logα)

test/layers/heteroconv.jl

@@ -148,4 +148,12 @@
         y = layers(hg, x); 
         @test size(y.A) == (2, 2) && size(y.B) == (2, 3)
     end
+
+    @testset "GATv2Conv" begin
+        x = (A = rand(Float32, 4, 2), B = rand(Float32, 4, 3))
+        layers = HeteroGraphConv((:A, :to, :B) => GATv2Conv(4 => 2),
+                                 (:B, :to, :A) => GATv2Conv(4 => 2));
+        y = layers(hg, x); 
+        @test size(y.A) == (2, 2) && size(y.B) == (2, 3)
+    end
 end