create GNNLux.jl package (#460)

* create GNNLux

* create GNNLux.jl

* fix ci

.github/workflows/test_GNNLux.yml

@@ -0,0 +1,48 @@
+name: GNNLux
+on:
+  pull_request:
+    branches:
+      - master
+  push:
+    branches:
+      - master
+jobs:
+  test:
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10' # Replace this with the minimum Julia version that your package supports.
+          # - '1' #  '1' will automatically expand to the latest stable 1.x release of Julia.
+          # - 'pre'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - name: Install Julia dependencies and run tests
+        shell: julia --project=monorepo {0}
+        run: |
+          using Pkg
+          # dev mono repo versions
+          pkg"registry up"
+          Pkg.update()
+          pkg"dev ./GNNGraphs ./GNNlib ./GNNLux"
+          Pkg.test("GNNLux"; coverage=true)
+      - uses: julia-actions/julia-processcoverage@v1
+        with:
+          # directories: ./GNNLux/src, ./GNNLux/ext
+          directories: ./GNNLux/src
+      - uses: codecov/codecov-action@v4
+        with:
+          files: lcov.info

GNNLux/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Carlo Lucibello <[email protected]> and contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

GNNLux/Project.toml

@@ -0,0 +1,35 @@
+name = "GNNLux"
+uuid = "e8545f4d-a905-48ac-a8c4-ca114b98986d"
+authors = ["Carlo Lucibello and contributors"]
+version = "0.1.0"
+
+[deps]
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
+GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+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", "ComponentArrays", "Functors", "LuxTestUtils", "ReTestItems", "StableRNGs", "Zygote"]

GNNLux/README.md

@@ -0,0 +1,2 @@
+# GNNLux.jl
+

GNNLux/src/GNNLux.jl

@@ -0,0 +1,15 @@
+module GNNLux
+using ConcreteStructs: @concrete
+using NNlib: NNlib
+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")
+export GraphConv
+
+end #module
+

GNNLux/src/layers/conv.jl

@@ -0,0 +1,93 @@
+
+@doc raw"""
+    GraphConv(in => out, σ=identity; aggr=+, bias=true, init=glorot_uniform)
+
+Graph convolution layer from Reference: [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://arxiv.org/abs/1810.02244).
+
+Performs:
+```math
+\mathbf{x}_i' = W_1 \mathbf{x}_i + \square_{j \in \mathcal{N}(i)} W_2 \mathbf{x}_j
+```
+
+where the aggregation type is selected by `aggr`.
+
+# Arguments
+
+- `in`: The dimension of input features.
+- `out`: The dimension of output features.
+- `σ`: Activation function.
+- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
+- `bias`: Add learnable bias.
+- `init`: Weights' initializer.
+
+# Examples
+
+```julia
+# create data
+s = [1,1,2,3]
+t = [2,3,1,1]
+in_channel = 3
+out_channel = 5
+g = GNNGraph(s, t)
+x = randn(Float32, 3, g.num_nodes)
+
+# create layer
+l = GraphConv(in_channel => out_channel, relu, bias = false, aggr = mean)
+
+# forward pass
+y = l(g, x)       
+```
+"""
+@concrete struct GraphConv <: AbstractExplicitLayer
+    in_dims::Int
+    out_dims::Int
+    use_bias::Bool
+    init_weight::Function
+    init_bias::Function
+    σ
+    aggr
+end
+
+
+function GraphConv(ch::Pair{Int, Int}, σ = identity;
+            aggr = +,
+            init_weight = glorot_uniform,
+            init_bias = zeros32, 
+            use_bias::Bool = true, 
+            allow_fast_activation::Bool = true)
+    in_dims, out_dims = ch
+    σ = allow_fast_activation ? NNlib.fast_act(σ) : σ
+    return GraphConv(in_dims, out_dims, use_bias, init_weight, init_bias, σ, aggr)
+end
+
+function LuxCore.initialparameters(rng::AbstractRNG, l::GraphConv)
+    weight1 = l.init_weight(rng, l.out_dims, l.in_dims)
+    weight2 = l.init_weight(rng, l.out_dims, l.in_dims)
+    if l.use_bias
+        bias = l.init_bias(rng, l.out_dims)
+    else
+        bias = false
+    end
+    return (; weight1, weight2, bias)
+end
+
+function LuxCore.parameterlength(l::GraphConv)
+    if l.use_bias
+        return 2 * l.in_dims * l.out_dims + l.out_dims
+    else
+        return 2 * l.in_dims * l.out_dims
+    end
+end
+
+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)
+    (l.σ == identity) || print(io, ", ", l.σ)
+    (l.aggr == +) || print(io, ", aggr=", l.aggr)
+    l.use_bias || print(io, ", use_bias=false")
+    print(io, ")")
+end
+
+(l::GraphConv)(g::GNNGraph, x, ps, st) = GNNlib.graph_conv(l, g, x, ps), st

GNNLux/test/layers/conv_tests.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/runtests.jl

@@ -0,0 +1,10 @@
+using Test
+using Lux
+using GNNLux
+using Random, Statistics
+
+using ReTestItems
+# using Pkg, Preferences, Test
+# using InteractiveUtils, Hwloc
+
+runtests(GNNLux)

GNNLux/test/shared_testsetup.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

GNNlib/src/layers/conv.jl

@@ -90,12 +90,15 @@ function  cheb_conv(c, g::GNNGraph, X::AbstractMatrix{T}) where {T}
     return Y .+ c.bias
 end
 
-function graph_conv(l, g::AbstractGNNGraph, x)
+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.σ.(l.weight1 * xi .+ l.weight2 * m .+ l.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)