Add SAGEConv support to HeteroGraphConv (#384)

* SAGEConv Hetero Layer

* without tests

* test doesnt work yet

* tests

* tests should work

* test update

* temporary testing fast

* final tests

---------

Co-authored-by: rbSparky <[email protected]>

Author: rbSparky

src/layers/conv.jl

@@ -801,10 +801,11 @@ function SAGEConv(ch::Pair{Int, Int}, σ = identity; aggr = mean,
     SAGEConv(W, b, σ, aggr)
 end
 
-function (l::SAGEConv)(g::GNNGraph, x::AbstractMatrix)
+function (l::SAGEConv)(g::AbstractGNNGraph, x)
     check_num_nodes(g, x)
-    m = propagate(copy_xj, g, l.aggr, xj = x)
-    x = l.σ.(l.weight * vcat(x, m) .+ l.bias)
+    xj, xi = expand_srcdst(g, x)
+    m = propagate(copy_xj, g, l.aggr, xj = xj)
+    x = l.σ.(l.weight * vcat(xi, m) .+ l.bias)
     return x
 end
 

test/layers/heteroconv.jl

@@ -109,4 +109,11 @@
         @test size(y.A) == (2,2) && size(y.B) == (2,3)
     end
 
+    @testset "SAGEConv" begin
+        x = (A = rand(Float32, 4, 2), B = rand(Float32, 4, 3))
+        layers = HeteroGraphConv((:A, :to, :B) => SAGEConv(4 => 2, relu, bias = false, aggr = +),
+                                 (:B, :to, :A) => SAGEConv(4 => 2, relu, bias = false, aggr = +));
+        y = layers(hg, x); 
+        @test size(y.A) == (2, 2) && size(y.B) == (2, 3)
+    end
 end