[GNNLux] Adding NNConv Layer

[rbSparky]

Adding conv layer according to #461

[rbSparky]

If some of the args are not needed for NNConv please let me know

[rbSparky]

Writing tests

[rbSparky]

dimension mismatch is in GNNLib implementation of nn conv:

function nn_conv(l, g::GNNGraph, x::AbstractMatrix, e)
    check_num_nodes(g, x)
    message = Fix1(nn_conv_message, l)
    m = propagate(message, g, l.aggr, xj = x, e = e)
    return l.σ.(l.weight * x .+ m .+ l.bias)
end

in the return statement

any suggestions or ideas as to why this could be happening?

[CarloLucibello]

This code works fine

using GraphNeuralNetworks, Flux

n_in = 3
n_in_edge = 10
n_out = 5

s = [1,1,2,3]
t = [2,3,1,1]
g = GNNGraph(s, t)

nn = Dense(n_in_edge => n_out * n_in)
l = NNConv(n_in => n_out, nn, tanh, bias = true, aggr = +)
x = randn(Float32, n_in, g.num_nodes)
e = randn(Float32, n_in_edge, g.num_edges)
y = l(g, x, e)  

Try to run the corresponding LuxGNN version and see if you get an error, we'll try to debug from there

[CarloLucibello]

any progress here?

[rbSparky]

any progress here?

yes got the same error

Got exception outside of a @test
  DimensionMismatch: A has dimensions (15,10) but B has dimensions (3,3)

[CarloLucibello]

Can you paste the code that errors?

[rbSparky]

Can you paste the code that errors?

i used the code that you gave for tests (in the test file) just without the bias, but that shouldnt cause an issue since issue is with multiplication
it errors while testing

[CarloLucibello]

as I said, can you translate this

using GraphNeuralNetworks, Flux

n_in = 3
n_in_edge = 10
n_out = 5

s = [1,1,2,3]
t = [2,3,1,1]
g = GNNGraph(s, t)

nn = Dense(n_in_edge => n_out * n_in)
l = NNConv(n_in => n_out, nn, tanh, bias = true, aggr = +)
x = randn(Float32, n_in, g.num_nodes)
e = randn(Float32, n_in_edge, g.num_edges)
y = l(g, x, e)  

to the corresponding GNNLux code and see what happens?

[rbSparky]

tried to run locally today but seems like there is an update to some package that leads to this error:
ERROR: LoadError: Failed to precompile MLDataDevices

while importing GNNlib (usual setup for the repo, activate, instantiate etc)

Could you check once as well?

[rbSparky]

Try to run the corresponding LuxGNN version

        n_in = 3
        n_in_edge = 10
        n_out = 5
        
        s = [1,1,2,3]
        t = [2,3,1,1]
        g2 = GNNGraph(s, t)
        
        nn = Dense(n_in_edge => n_out * n_in)
        l = NNConv(n_in => n_out, nn, tanh, aggr = +)
        x = randn(Float32, n_in, g2.num_nodes)
        e = randn(Float32, n_in_edge, g2.num_edges)

        ps = LuxCore.initialparameters(rng, l)
        st = LuxCore.initialstates(rng, l)
        
        y = l(g2, x, e, ps, st)

added this to the test file, there we can see what errors

[rbSparky]

error:

  DimensionMismatch: A has dimensions (15,10) but B has dimensions (3,3)

[rbSparky]

Any suggestions?

[rbSparky]

getting the same dimension error

[CarloLucibello]

The problem is that ps.weight corresponded to the internal weights of the nn in Lux < v1.0 due LuxDL/Lux.jl#795.
So you were passing to GNNLib.nn_conv the weights of the wrong shape to use in this line

In order to fix the layer you need to do two things:

• rebase this branch on top of master, so that now we work with Lux >= v1.0 and the weights in the internal nn are found in ps.nn.weight (this was a breaking change in Lux' container layer)
• Use this definition for the layer (didn't fully test it though)

@concrete struct NNConv <: GNNContainerLayer{(:nn,)}
    nn <: AbstractLuxLayer    
    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)
    in_dims, out_dims = ch
    σ = NNlib.fast_act(σ)
    return NNConv(nn, aggr, in_dims, out_dims, use_bias, init_weight, init_bias, σ)
end

function LuxCore.initialparameters(rng::AbstractRNG, l::NNConv)
    weight = l.init_weight(rng, l.out_dims, l.in_dims)
    ps = (; nn = LuxCore.initialparameters(rng, l.nn), weight)
    if l.use_bias
        ps = (; ps..., bias = l.init_bias(rng, l.out_dims))
    end
    return ps
end

function LuxCore.initialstates(rng::AbstractRNG, l::NNConv)
    return (; nn = LuxCore.initialstates(rng, l.nn))
end

function LuxCore.parameterlength(l::NNConv)
    n = parameterlength(l.nn) + l.in_dims * l.out_dims
    if l.use_bias
        n += l.out_dims
    end
    return n
end

LuxCore.statelength(l::NNConv) = statelength(l.nn)

function (l::NNConv)(g, x, e, ps, st)
    nn = StatefulLuxLayer{true}(l.nn, ps.nn, st.nn)
    m = (; nn, l.aggr, ps.weight, bias = _getbias(ps), l.σ)
    y = GNNlib.nn_conv(m, g, x, e)
    stnew = _getstate(nn)
    return y, stnew
end

[CarloLucibello]

Since in this branch there are so many commits and a merge rebase could be hard. Close this and open a new PR if needed.

[rbSparky]

rebasing and making new PR

[rbSparky]

#491