Updated self loop logic and tests

Author: Amogh

GNNLux/src/layers/conv.jl

@@ -422,8 +422,7 @@ function Base.show(io::IO, l::AGNNConv)
 end
 
 function (l::AGNNConv)(g::AbstractGNNGraph, x, ps, st)
-    β = l.trainable ? ps.β : l.init_beta
-    m = (;  β, l.add_self_loops)
+    β = l.trainable ? ps.β[1] : l.init_beta[1]  
     return GNNlib.agnn_conv(g, x, β; self_loops=l.add_self_loops), st
 end
 

GNNlib/src/layers/conv.jl

@@ -336,26 +336,22 @@ end
 
 function agnn_conv(g::AbstractGNNGraph, x, β; self_loops=true)
 
-    if self_loops && !(g.num_nodes isa Integer)
+    if self_loops && (!(g isa GNNHeteroGraph) || all(et -> et[1] == et[3], g.etypes))
         g = add_self_loops(g)
     end
 
     xj, xi = expand_srcdst(g, x)
-    T = eltype(xj)
-    ϵ = T(1e-9)
-
-    xi_norm = xi ./ sqrt.(sum(abs2, xi, dims=1) .+ ϵ)
-    xj_norm = xj ./ sqrt.(sum(abs2, xj, dims=1) .+ ϵ)
-
-    s, d = edge_index(g)
 
-    cos_dist = sum(xi_norm[:, d] .* xj_norm[:, s], dims=1)
+    xi_norm = xi ./ sqrt.(sum(abs2, xi, dims=1) .+ eps(eltype(xi)))
+    xj_norm = xj ./ sqrt.(sum(abs2, xj, dims=1) .+ eps(eltype(xj)))
 
+    cos_dist = apply_edges(xi_dot_xj, g, xi=xi_norm, xj=xj_norm)
+    
     α = softmax_edge_neighbors(g, β .* cos_dist)
 
-    m = α .* xj[:, s]
-    
-    return GNNlib.aggregate_neighbors(g, +, m)
+    return propagate(g, +; xj=xj, e=α) do xi_i, xj_j, α_e
+        α_e .* xj_j
+    end
 end
 
 ####################### MegNetConv ######################################

GraphNeuralNetworks/test/layers/conv.jl

@@ -397,7 +397,7 @@ end
 # end
 
 @testitem "AGNNConv" setup=[TolSnippet, TestModule] begin
-    using .TestModule
+    using .TestModule    
     @testset "Initialization & Basic Forward" begin
         l = AGNNConv(init_beta=2.0f0, trainable=true)
         @test l.β[1] == 2.0f0 
@@ -406,25 +406,28 @@ end
             @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
+    end    
+    @testset "Heterogeneous Graph Support" begin
+        l = AGNNConv(add_self_loops=false)  
+        hg = rand_bipartite_heterograph((10, 15), 20)
+        x = (A = randn(Float32, D_IN, 10), B = randn(Float32, D_IN, 15))        
+        hetero_layer = HeteroGraphConv(
+            (:A, :to, :B) => l,  
+            (:B, :to, :A) => l;  
+            aggr = +
+        )        
+        y = hetero_layer(hg, x)
+        @test size(y.A) == (D_IN, 10)
+        @test size(y.B) == (D_IN, 15)
+    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
+        x_dead[:, 1] .= 0.0f0  
         y = l(g, x_dead)
         @test !any(isnan.(y)) 
+        @test !any(isinf.(y))
     end
 end
 
@@ -435,7 +438,14 @@ end
         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
+    l_bip = AGNNConv(add_self_loops=false)
+    s = [1, 1, 2, 3]
+    t = [1, 2, 1, 2]
+    g = GNNGraph((s, t)) |> gpu
+    x = (randn(Float32, D_IN, 3) |> gpu, randn(Float32, D_IN, 2) |> gpu)
+    y = l_bip(g, x)
+    @test size(y) == (D_IN, 2)
 end
 
 @testitem "MEGNetConv" setup=[TolSnippet, TestModule] begin