Added support for HeteroGraphConv in AGNNConv

Author: Amogh

GNNLux/src/layers/conv.jl

@@ -421,10 +421,10 @@ function Base.show(io::IO, l::AGNNConv)
     print(io, ")")
 end
 
-function (l::AGNNConv)(g, x::AbstractMatrix, ps, st)
+function (l::AGNNConv)(g::AbstractGNNGraph, x, ps, st)
     β = l.trainable ? ps.β : l.init_beta
     m = (;  β, l.add_self_loops)
-    return GNNlib.agnn_conv(m, g, x), st
+    return GNNlib.agnn_conv(g, x, β; self_loops=l.add_self_loops), st
 end
 
 @doc raw"""

GNNlib/Project.toml

@@ -1,7 +1,7 @@
 name = "GNNlib"
 uuid = "a6a84749-d869-43f8-aacc-be26a1996e48"
 authors = ["Carlo Lucibello and contributors"]
-version = "1.2.0"
+version = "1.3.0-DEV"
 
 [workspace]
 projects = ["test", "docs", "../GNNGraphs"]

GNNlib/src/layers/conv.jl

@@ -334,21 +334,28 @@ end
 
 ####################### AGNNConv ######################################
 
-function agnn_conv(l, g::GNNGraph, x::AbstractMatrix)
-    check_num_nodes(g, x)
-    if l.add_self_loops
+function agnn_conv(g::AbstractGNNGraph, x, β; self_loops=true)
+    
+    if self_loops && !(g.num_nodes isa Integer)
         g = add_self_loops(g)
     end
 
-    xn = x ./ sqrt.(sum(x .^ 2, dims = 1))
-    cos_dist = apply_edges(xi_dot_xj, g, xi = xn, xj = xn)
-    α = softmax_edge_neighbors(g, l.β .* cos_dist)
+    xj, xi = expand_srcdst(g, x)
+    T = eltype(xj)
+    ϵ = T(1e-9)
 
-    x = propagate(g, +; xj = x, e = α) do xi, xj, α
-        α .* xj 
-    end
+    xi_norm = xi ./ sqrt.(sum(abs2, xi, dims=1) .+ ϵ)
+    xj_norm = xj ./ sqrt.(sum(abs2, xj, dims=1) .+ ϵ)
 
-    return x
+    s, d = edge_index(g)
+    
+    cos_dist = sum(xi_norm[:, d] .* xj_norm[:, s], dims=1)
+
+    α = softmax_edge_neighbors(g, β .* cos_dist)
+
+    m = α .* xj[:, s]
+    
+    return GNNlib.aggregate_neighbors(g, +, m)
 end
 
 ####################### MegNetConv ######################################

GraphNeuralNetworks/Project.toml

@@ -1,12 +1,10 @@
 name = "GraphNeuralNetworks"
 uuid = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 authors = ["Carlo Lucibello and contributors"]
-version = "1.1.0"
-
-[workspace]
-projects = ["test", "docs"]
+version = "1.2.0-DEV"
 
 [deps]
+AutoStructs = "2e0df379-9877-4907-ab94-cd881f8d985b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
@@ -21,6 +19,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+AutoStructs = "0.1.1"
 ChainRulesCore = "1"
 ConcreteStructs = "0.2.3"
 Flux = "0.16.0"
@@ -34,3 +33,6 @@ Random = "1"
 Reexport = "1"
 Statistics = "1"
 julia = "1.10"
+
+[workspace]
+projects = ["test", "docs"]

GraphNeuralNetworks/src/layers/conv.jl

@@ -4,6 +4,8 @@
 # the implementation of a single layer, 
 # but it is done for GraphNeuralNetworks.jl and GNNLux.jl to be able to share the same code.
 
+using Flux, GraphNeuralNetworks, AutoStructs
+
 @doc raw"""
     GCNConv(in => out, σ=identity; [bias, init, add_self_loops, use_edge_weight])
 
@@ -999,7 +1001,9 @@ function AGNNConv(; init_beta = 1.0f0, add_self_loops = true, trainable = true)
     AGNNConv([init_beta], add_self_loops, trainable)
 end
 
-(l::AGNNConv)(g, x) = GNNlib.agnn_conv(l, g, x)
+function (l::AGNNConv)(g::AbstractGNNGraph, x)
+    return GNNlib.agnn_conv(g, x, l.β[1]; self_loops=l.add_self_loops)
+end
 
 @doc raw"""
     MEGNetConv(ϕe, ϕv; aggr=mean)

GraphNeuralNetworks/test/layers/conv.jl

@@ -398,29 +398,42 @@ end
 
 @testitem "AGNNConv" setup=[TolSnippet, TestModule] begin
     using .TestModule
-    l = AGNNConv(trainable=false, add_self_loops=false)
-    @test l.β == [1.0f0]
-    @test l.add_self_loops == false
-    @test l.trainable == false
-    Flux.trainable(l) == (;)
-
-    l = AGNNConv(init_beta=2.0f0)
-    @test l.β == [2.0f0]
-    @test l.add_self_loops == true
-    @test l.trainable == true 
-    Flux.trainable(l) == (; β = [1f0])
-    for g in TEST_GRAPHS
-        @test size(l(g, g.x)) == (D_IN, g.num_nodes)
-        test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_mooncake = TEST_MOONCAKE)
+    @testset "Initialization & Basic Forward" begin
+        l = AGNNConv(init_beta=2.0f0, trainable=true)
+        @test l.β[1] == 2.0f0 
+        @test l.trainable == true 
+        for g in TEST_GRAPHS
+            @test size(l(g, g.x)) == (D_IN, g.num_nodes) 
+            test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_mooncake = TEST_MOONCAKE) 
+        end
+    end
+    @testset "Bipartite Support" begin
+        l = AGNNConv(add_self_loops=false)
+        in_channel = 16
+        s = [rand(1:5, 14)..., 5] 
+        d = [rand(1:8, 14)..., 8]
+        g = GNNGraph((s, d))
+        x = (randn(Float32, in_channel, 5), randn(Float32, in_channel, 8))
+        y = l(g, x) 
+        @test size(y) == (in_channel, 8) 
+        test_gradients(l, g, x, rtol = RTOL_HIGH, test_mooncake = TEST_MOONCAKE)
+    end
+    @testset "Stability (Epsilon)" begin
+        l = AGNNConv()
+        g = TEST_GRAPHS[1]
+        x_dead = randn(Float32, D_IN, g.num_nodes)
+        x_dead[:, 1] .= 0.0f0
+        y = l(g, x_dead)
+        @test !any(isnan.(y)) 
     end
 end
 
 @testitem "AGNNConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
     using .TestModule
-    l = AGNNConv(trainable=false, add_self_loops=false)
+    l = AGNNConv()
     for g in TEST_GRAPHS
-        g.graph isa AbstractSparseMatrix && continue
-        @test size(l(g, g.x)) == (D_IN, g.num_nodes)
+        g.graph isa AbstractSparseMatrix && continue 
+        @test size(l(g, g.x)) == (D_IN, g.num_nodes) 
         test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
     end   
 end

GraphNeuralNetworks/test/layers/heteroconv.jl

@@ -158,4 +158,14 @@
         y = layers(g, x); 
         @test size(y.A) == (2,2) && size(y.B) == (2,3)
     end
+
+    @testset "AGNNConv" begin
+        x = (A = rand(Float32, 4, 2), B = rand(Float32, 4, 3))
+        
+        layers = HeteroGraphConv((:A, :to, :B) => AGNNConv(add_self_loops=false),
+                                 (:B, :to, :A) => AGNNConv(add_self_loops=false))
+        
+        y = layers(hg, x)
+        @test size(y.A) == (4, 2) && size(y.B) == (4, 3)
+    end
 end