Adding NNConv support for HeteroGraphConv

Author: Amogh

GNNGraphs/src/gnnheterograph/query.jl

@@ -89,3 +89,13 @@ function graph_indicator(g::GNNHeteroGraph, node_t::Symbol)
     end
     return gi
 end
+
+edge_features(g::GNNHeteroGraph, edge_t::EType) = begin
+    ds == g.edata[edge_t]
+    isempty(ds) ? nothing : first(values(ds))
+end
+
+edge_features(g::GNNHeteroGraph) = begin
+    ds = only(values(g.edata))
+    isempty(ds) ? nothing : first(values(ds))
+end

GNNlib/src/layers/conv.jl

@@ -257,11 +257,12 @@ end
 
 ####################### NNConv ######################################
 
-function nn_conv(l, g::GNNGraph, x::AbstractMatrix, e)
+function nn_conv(l, g::AbstractGNNGraph, x, e)
     check_num_nodes(g, x)
+    xj, xi = expand_srcdst(g, x)
     message = Fix1(nn_conv_message, l)
-    m = propagate(message, g, l.aggr, xj = x, e = e)
-    return l.σ.(l.weight * x .+ m .+ l.bias)
+    m = propagate(message, g, l.aggr, xj = xj, e = e)
+    return l.σ.(l.weight * xi .+ m .+ l.bias)
 end
 
 function nn_conv_message(l, xi, xj, e)

GraphNeuralNetworks/src/layers/conv.jl

@@ -718,6 +718,8 @@ end
 
 (l::NNConv)(g, x, e) = GNNlib.nn_conv(l, g, x, e)
 
+(l::NNConv)(g, x) = GNNlib.nn_conv(l, g, x, edge_features(g))
+
 (l::NNConv)(g::GNNGraph) = GNNGraph(g, ndata = l(g, node_features(g), edge_features(g)))
 
 function Base.show(io::IO, l::NNConv)

GraphNeuralNetworks/test/layers/heteroconv.jl

@@ -155,7 +155,22 @@
         x = (A = rand(Float32, 4,2), B = rand(Float32, 4, 3))
         layers = HeteroGraphConv( (:A, :to, :B) => GCNConv(4 => 2, tanh),
                                     (:B, :to, :A) => GCNConv(4 => 2, tanh));
-        y = layers(g, x); 
+        y = layers(g, x);
         @test size(y.A) == (2,2) && size(y.B) == (2,3)
     end
+
+    @testset "NNConv" begin
+        nA, nB = 2, 3
+        nedges = 6
+        g = rand_bipartite_heterograph((nA, nB), nedges;
+            edata = Dict((:A, :to, :B) => (e = rand(Float32, 5, nedges),),
+                        (:B, :to, :A) => (e = rand(Float32, 5, nedges),))
+        )
+        x = (A = rand(Float32, 4, nA), B = rand(Float32, 4, nB))
+        nn = Dense(5 => 4 * 2)
+        layers = HeteroGraphConv((:A, :to, :B) => NNConv(4 => 2, nn, tanh),
+                                 (:B, :to, :A) => NNConv(4 => 2, nn, tanh))
+        y = layers(g, x)
+        @test size(y.A) == (2, nA) && size(y.B) == (2, nB)
+    end
 end