create GNNLux.jl

Author: CarloLucibello

GNNLux/Project.toml

@@ -12,14 +12,24 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-
-[extensions]
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+ConcreteStructs = "0.2.3"
+Lux = "0.5.61"
+LuxCore = "0.1.20"
+NNlib = "0.9.21"
+Reexport = "1.2"
 julia = "1.10"
 
 [extras]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test"]
+test = ["Test", "ComponentArrays", "Functors", "LuxTestUtils", "ReTestItems", "StableRNGs", "Zygote"]

GNNLux/src/GNNLux.jl

@@ -5,7 +5,7 @@ using LuxCore: LuxCore, AbstractExplicitLayer
 using Lux: glorot_uniform, zeros32
 using Reexport: @reexport
 using Random: AbstractRNG
-
+using GNNlib: GNNlib
 @reexport using GNNGraphs
 
 include("layers/conv.jl")

GNNLux/src/layers/conv.jl

@@ -79,8 +79,8 @@ function LuxCore.parameterlength(l::GraphConv)
     end
 end
 
-statelength(d::GraphConv) = 0
-outputsize(d::GraphConv) = (d.out_dims,)
+LuxCore.statelength(d::GraphConv) = 0
+LuxCore.outputsize(d::GraphConv) = (d.out_dims,)
 
 function Base.show(io::IO, l::GraphConv)
     print(io, "GraphConv(", l.in_dims, " => ", l.out_dims)

GNNLux/test/layers/conv.jl

@@ -0,0 +1,19 @@
+@testitem "layers/conv" setup=[SharedTestSetup] begin
+    rng = StableRNG(1234)
+    g = rand_graph(10, 30, seed=1234)
+    x = randn(rng, Float32, 3, 10)
+
+    @testset "GraphConv" begin
+        l = GraphConv(3 => 5, relu)
+        ps = Lux.initialparameters(rng, l)
+        st = Lux.initialstates(rng, l)
+        @test Lux.parameterlength(l) == Lux.parameterlength(ps)
+        @test Lux.statelength(l) == Lux.statelength(st)
+
+        y, _ = l(g, x, ps, st)
+        @test Lux.outputsize(l) == (5,)
+        @test size(y) == (5, 10)
+        loss = (x, ps) -> sum(first(l(g, x, ps, st)))
+        @eval @test_gradients $loss $x $ps atol=1.0f-3 rtol=1.0f-3
+    end
+end

GNNLux/test/layers/conv_tests.jl

@@ -3,18 +3,8 @@ using Lux
 using GNNLux
 using Random, Statistics
 
+using ReTestItems
+# using Pkg, Preferences, Test
+# using InteractiveUtils, Hwloc
 
-tests = [
-    # "utils",
-    # "msgpass",
-    # "layers/basic",
-    "layers/conv",
-    # "layers/heteroconv",
-    # "layers/temporalconv",
-    # "layers/pool",
-    # "examples/node_classification_cora",
-]
-
-@testset "$t" for t in tests
-    include("$t.jl")
-end
+runtests(GNNLux)

GNNLux/test/runtests.jl

@@ -0,0 +1,23 @@
+@testsetup module SharedTestSetup
+
+import Reexport: @reexport
+
+@reexport using Lux, Functors
+@reexport using ComponentArrays, LuxCore, LuxTestUtils, Random, StableRNGs, Test,
+                Zygote, Statistics
+@reexport using LuxTestUtils: @jet, @test_gradients, check_approx
+
+# Some Helper Functions
+function get_default_rng(mode::String)
+    dev = mode == "cpu" ? LuxCPUDevice() :
+          mode == "cuda" ? LuxCUDADevice() : mode == "amdgpu" ? LuxAMDGPUDevice() : nothing
+    rng = default_device_rng(dev)
+    return rng isa TaskLocalRNG ? copy(rng) : deepcopy(rng)
+end
+
+export get_default_rng
+
+# export BACKEND_GROUP, MODES, cpu_testing, cuda_testing, amdgpu_testing, get_default_rng,
+#        StableRNG, maybe_rewrite_to_crosscor
+
+end

GNNLux/test/shared_testsetup.jl

@@ -94,8 +94,11 @@ function graph_conv(l, g::AbstractGNNGraph, x, ps)
     check_num_nodes(g, x)
     xj, xi = expand_srcdst(g, x)
     m = propagate(copy_xj, g, l.aggr, xj = xj)
-    x = l.σ.(ps.weight1 * xi .+ ps.weight2 * m .+ ps.bias)
-    return x
+    x = ps.weight1 * xi .+ ps.weight2 * m
+    if l.use_bias
+        x = x .+ ps.bias
+    end
+    return l.σ.(x)
 end
 
 function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = nothing)