[GNNLux] Adding NNConv Layer

GNNLux/Project.toml

@@ -4,15 +4,29 @@ authors = ["Carlo Lucibello and contributors"]
 version = "0.1.0"
 
 [deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
 GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
+GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ConcreteStructs = "0.2.3"

GNNLux/src/GNNLux.jl

@@ -31,8 +31,8 @@ export AGNNConv,
        GINConv,
        # GMMConv,
        GraphConv,
+       NNConv,
        MEGNetConv,
-       # NNConv,
        # ResGatedGraphConv,
        # SAGEConv,
        SGConv

GNNLux/src/layers/conv.jl

@@ -629,6 +629,44 @@ function Base.show(io::IO, l::GINConv)
     print(io, ")")
 end
 
+@concrete struct NNConv <: GNNContainerLayer{(:nn,)}
+    nn <: AbstractExplicitLayer    
+    aggr
+    in_dims::Int
+    out_dims::Int
+    use_bias::Bool
+    init_weight
+    init_bias
+    σ
+end
+
+function NNConv(ch::Pair{Int, Int}, nn, σ = identity; 
+                aggr = +, 
+                init_bias = zeros32,
+                use_bias::Bool = true,
+                init_weight = glorot_uniform,
+                allow_fast_activation::Bool = true)
+    in_dims, out_dims = ch
+    σ = allow_fast_activation ? NNlib.fast_act(σ) : σ
+    return NNConv(nn, aggr, in_dims, out_dims, use_bias, init_weight, init_bias, σ)
+end
+
+function (l::NNConv)(g, x, edge_weight, ps, st)
+    nn = StatefulLuxLayer{true}(l.nn, ps, st)
+
+    m = (; nn, l.aggr, ps.weight, bias = _getbias(ps), l.σ)
+    y = GNNlib.nn_conv(m, g, x, edge_weight)
+    stnew = _getstate(nn)
+    return y, stnew
+end
+
+function Base.show(io::IO, l::NNConv)
+    print(io, "NNConv($(l.nn)")
+    l.σ == identity || print(io, ", ", l.σ)
+    l.use_bias || print(io, ", use_bias=false")
+    print(io, ")")
+end
+
 @concrete struct MEGNetConv{TE, TV, A} <: GNNContainerLayer{(:ϕe, :ϕv)}
     in_dims::Int
     out_dims::Int
@@ -669,4 +707,4 @@ function Base.show(io::IO, l::MEGNetConv)
     nout = l.out_dims
     print(io, "MEGNetConv(", nin, " => ", nout)
     print(io, ")")
-end
+end

GNNLux/test/layers/conv_tests.jl

@@ -1,10 +1,13 @@
 @testitem "layers/conv" setup=[SharedTestSetup] begin
     rng = StableRNG(1234)
+    edim = 10
     g = rand_graph(rng, 10, 40)
     in_dims = 3
     out_dims = 5
     x = randn(rng, Float32, in_dims, 10)
 
+    g2 = GNNGraph(g, edata = rand(Float32, edim, g.num_edges)) 
+
     @testset "GCNConv" begin
         l = GCNConv(in_dims => out_dims, tanh)
         test_lux_layer(rng, l, g, x, outputsize=(out_dims,))
@@ -94,6 +97,13 @@
         test_lux_layer(rng, l, g, x, sizey=(out_dims,g.num_nodes), container=true)
     end
 
+    @testset "NNConv" begin
+        edim = 10
+        nn = Dense(edim, in_dims * out_dims)        
+        l = NNConv(in_dims => out_dims, nn, tanh, aggr = +)
+        test_lux_layer(rng, l, g2, x, sizey=(out_dims, g2.num_nodes), container=true, edge_weight=g2.edata.e) 
+    end          
+
     @testset "MEGNetConv" begin
         l = MEGNetConv(in_dims => out_dims)
 

GNNLux/test/shared_testsetup.jl

@@ -14,7 +14,7 @@ export test_lux_layer
 
 function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x; 
             outputsize=nothing, sizey=nothing, container=false,
-            atol=1.0f-2, rtol=1.0f-2)
+            atol=1.0f-2, rtol=1.0f-2, edge_weight=nothing) 
 
     if container
         @test l isa GNNContainerLayer
@@ -26,8 +26,13 @@ function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x;
     st = LuxCore.initialstates(rng, l)
     @test LuxCore.parameterlength(l) == LuxCore.parameterlength(ps)
     @test LuxCore.statelength(l) == LuxCore.statelength(st)
-
-    y, st′ = l(g, x, ps, st)
+
+    if edge_weight !== nothing
+        y, st′ = l(g, x, edge_weight, ps, st)
+    else          
+        y, st′ = l(g, x, ps, st)
+    end
+
     @test eltype(y) == eltype(x)
     if outputsize !== nothing
         @test LuxCore.outputsize(l) == outputsize
@@ -42,4 +47,4 @@ function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x;
     test_gradients(loss, x, ps; atol, rtol, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
 end
 
-end
+end