feat: Add EdgeConv support for HeteroGraphConv (#364)

* add edgeconv heterograph support

* add test & resolve conflicts

* cleanup

* Update test/layers/heteroconv.jl

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

* cleanup after rebase

---------

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

Author: askorupka

src/layers/conv.jl

@@ -637,10 +637,13 @@ end
 
 EdgeConv(nn; aggr = max) = EdgeConv(nn, aggr)
 
-function (l::EdgeConv)(g::GNNGraph, x::AbstractMatrix)
+function (l::EdgeConv)(g::AbstractGNNGraph, x)
     check_num_nodes(g, x)
+    xj, xi = expand_srcdst(g, x)
+
     message(l, xi, xj, e) = l.nn(vcat(xi, xj .- xi))
-    x = propagate(message, g, l.aggr, l, xi = x, xj = x)
+
+    x = propagate(message, g, l.aggr, l, xi = xi, xj = xj, e = nothing)
     return x
 end
 

test/layers/heteroconv.jl

@@ -100,4 +100,14 @@
         y = layers(g, x); 
         @test size(y.A) == (2,2) && size(y.B) == (2,3)
     end
+
+    @testset "EdgeConv" begin
+        g = rand_bipartite_heterograph((2,3), 6)
+        x = (A = rand(Float32, 4,2), B = rand(Float32, 4, 3))
+        layers = HeteroGraphConv( (:A, :to, :B) => EdgeConv(Dense(2 * 4, 2), aggr = +),
+                                    (:B, :to, :A) => EdgeConv(Dense(2 * 4, 2), aggr = +));
+        y = layers(g, x); 
+        @test size(y.A) == (2,2) && size(y.B) == (2,3)
+    end
+
 end