[GNNLux] more layers pt. 3 (#471)

* more layer

more layers

stuff

* fixes

GNNLux/src/GNNLux.jl

@@ -1,8 +1,11 @@
 module GNNLux
 using ConcreteStructs: @concrete
 using NNlib: NNlib, sigmoid, relu, swish
-using LuxCore: LuxCore, AbstractExplicitLayer, AbstractExplicitContainerLayer
-using Lux: Lux, Chain, Dense, glorot_uniform, zeros32, StatefulLuxLayer
+using LuxCore: LuxCore, AbstractExplicitLayer, AbstractExplicitContainerLayer, parameterlength, statelength, outputsize, 
+ initialparameters, initialstates, parameterlength, statelength
+using Lux: Lux, Chain, Dense, GRUCell,
+ glorot_uniform, zeros32, 
+ StatefulLuxLayer
 using Reexport: @reexport
 using Random: AbstractRNG
 using GNNlib: GNNlib
@@ -22,9 +25,9 @@ export AGNNConv,
  DConv,
  GATConv,
  GATv2Conv,
- # GatedGraphConv,
+ GatedGraphConv,
  GCNConv,
- # GINConv,
+ GINConv,
  # GMMConv,
  GraphConv,
  # MEGNetConv,

GNNLux/src/layers/conv.jl

@@ -38,7 +38,6 @@ function LuxCore.initialparameters(rng::AbstractRNG, l::GCNConv)
 end
 
 LuxCore.parameterlength(l::GCNConv) = l.use_bias ? l.in_dims * l.out_dims + l.out_dims : l.in_dims * l.out_dims
-LuxCore.statelength(d::GCNConv) = 0
 LuxCore.outputsize(d::GCNConv) = (d.out_dims,)
 
 function Base.show(io::IO, l::GCNConv)
@@ -549,7 +548,6 @@ function LuxCore.initialparameters(rng::AbstractRNG, l::SGConv)
 end
 
 LuxCore.parameterlength(l::SGConv) = l.use_bias ? l.in_dims * l.out_dims + l.out_dims : l.in_dims * l.out_dims
-LuxCore.statelength(d::SGConv) = 0
 LuxCore.outputsize(d::SGConv) = (d.out_dims,)
 
 function Base.show(io::IO, l::SGConv)
@@ -561,14 +559,72 @@ function Base.show(io::IO, l::SGConv)
  print(io, ")")
 end
 
-(l::SGConv)(g, x, ps, st; conv_weight=nothing, edge_weight=nothing) = 
- l(g, x, edge_weight, ps, st; conv_weight)
-
-function (l::SGConv)(g, x, edge_weight, ps, st; 
- conv_weight=nothing, )
+(l::SGConv)(g, x, ps, st) = l(g, x, nothing, ps, st)
 
+function (l::SGConv)(g, x, edge_weight, ps, st)
  m = (; ps.weight, bias = _getbias(ps), 
  l.add_self_loops, l.use_edge_weight, l.k)
  y = GNNlib.sg_conv(m, g, x, edge_weight)
  return y, st
-end
+end
+
+@concrete struct GatedGraphConv <: GNNLayer
+ gru
+ init_weight
+ dims::Int
+ num_layers::Int
+ aggr
+end
+
+
+function GatedGraphConv(dims::Int, num_layers::Int; 
+ aggr = +, init_weight = glorot_uniform)
+ gru = GRUCell(dims => dims)
+ return GatedGraphConv(gru, init_weight, dims, num_layers, aggr)
+end
+
+LuxCore.outputsize(l::GatedGraphConv) = (l.dims,)
+
+function LuxCore.initialparameters(rng::AbstractRNG, l::GatedGraphConv)
+ gru = LuxCore.initialparameters(rng, l.gru)
+ weight = l.init_weight(rng, l.dims, l.dims, l.num_layers)
+ return (; gru, weight)
+end
+
+LuxCore.parameterlength(l::GatedGraphConv) = parameterlength(l.gru) + l.dims^2*l.num_layers
+
+
+function (l::GatedGraphConv)(g, x, ps, st)
+ gru = StatefulLuxLayer{true}(l.gru, ps.gru, _getstate(st, :gru))
+ fgru = (h, x) -> gru((x, (h,))) # make the forward compatible with Flux.GRUCell style
+ m = (; gru=fgru, ps.weight, l.num_layers, l.aggr, l.dims)
+ return GNNlib.gated_graph_conv(m, g, x), st
+end
+
+function Base.show(io::IO, l::GatedGraphConv)
+ print(io, "GatedGraphConv($(l.dims), $(l.num_layers)")
+ print(io, ", aggr=", l.aggr)
+ print(io, ")")
+end
+
+@concrete struct GINConv <: GNNContainerLayer{(:nn,)}
+ nn <: AbstractExplicitLayer
+ ϵ <: Real
+ aggr
+end
+
+GINConv(nn, ϵ; aggr = +) = GINConv(nn, ϵ, aggr)
+
+function (l::GINConv)(g, x, ps, st)
+ nn = StatefulLuxLayer{true}(l.nn, ps, st)
+ m = (; nn, l.ϵ, l.aggr)
+ y = GNNlib.gin_conv(m, g, x)
+ stnew = _getstate(nn)
+ return y, stnew
+end
+
+function Base.show(io::IO, l::GINConv)
+ print(io, "GINConv($(l.nn)")
+ print(io, ", $(l.ϵ)")
+ print(io, ")")
+end

GNNLux/test/layers/conv_tests.jl

@@ -82,5 +82,15 @@
  l = SGConv(in_dims => out_dims, 2)
  test_lux_layer(rng, l, g, x, outputsize=(out_dims,))
  end
-end
 
+ @testset "GatedGraphConv" begin
+ l = GatedGraphConv(in_dims, 3)
+ test_lux_layer(rng, l, g, x, outputsize=(in_dims,))
+ end
+
+ @testset "GINConv" begin
+ nn = Chain(Dense(in_dims => out_dims, relu), Dense(out_dims => out_dims))
+ l = GINConv(nn, 0.5)
+ test_lux_layer(rng, l, g, x, sizey=(out_dims,g.num_nodes), container=true)
+ end
+end

GNNLux/test/shared_testsetup.jl

@@ -28,6 +28,7 @@ function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x;
  @test LuxCore.statelength(l) == LuxCore.statelength(st)
 
  y, st′ = l(g, x, ps, st)
+ @test eltype(y) == eltype(x)
  if outputsize !== nothing
  @test LuxCore.outputsize(l) == outputsize
  end

GNNlib/src/layers/conv.jl

@@ -28,7 +28,7 @@ function gcn_conv(l, g::AbstractGNNGraph, x, edge_weight::EW, norm_fn::F, conv_w
  if edge_weight !== nothing
  # Pad weights with ones
  # TODO for ADJMAT_T the new edges are not generally at the end
- edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
+ edge_weight = [edge_weight; ones_like(edge_weight, g.num_nodes)]
  @assert length(edge_weight) == g.num_edges
  end
  end
@@ -215,23 +215,22 @@ end
 
 ####################### GatedGraphConv ######################################
 
-# TODO PIRACY! remove after https://github.com/JuliaDiff/ChainRules.jl/pull/521
-@non_differentiable fill!(x...)
-
-function gated_graph_conv(l, g::GNNGraph, H::AbstractMatrix{S}) where {S <: Real}
- check_num_nodes(g, H)
- m, n = size(H)
- @assert (m<=l.out_ch) "number of input features must less or equals to output features."
- if m < l.out_ch
- Hpad = similar(H, S, l.out_ch - m, n)
- H = vcat(H, fill!(Hpad, 0))
+function gated_graph_conv(l, g::GNNGraph, x::AbstractMatrix)
+ check_num_nodes(g, x)
+ m, n = size(x)
+ @assert m <= l.dims "number of input features must be less or equal to output features."
+ if m < l.dims
+ xpad = zeros_like(x, (l.dims - m, n))
+ x = vcat(x, xpad)
  end
+ h = x
  for i in 1:(l.num_layers)
- M = view(l.weight, :, :, i) * H
- M = propagate(copy_xj, g, l.aggr; xj = M)
- H, _ = l.gru(H, M)
+ m = view(l.weight, :, :, i) * h
+ m = propagate(copy_xj, g, l.aggr; xj = m)
+ # in gru forward, hidden state is first argument, input is second
+ h, _ = l.gru(h, m)
  end
- return H
+ return h
 end
 
 ####################### EdgeConv ######################################
@@ -419,7 +418,7 @@ function sgc_conv(l, g::GNNGraph, x::AbstractMatrix{T},
  if l.add_self_loops
  g = add_self_loops(g)
  if edge_weight !== nothing
- edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
+ edge_weight = [edge_weight; onse_like(edge_weight, g.num_nodes)]
  @assert length(edge_weight) == g.num_edges
  end
  end
@@ -512,7 +511,7 @@ function sg_conv(l, g::GNNGraph, x::AbstractMatrix{T},
  if l.add_self_loops
  g = add_self_loops(g)
  if edge_weight !== nothing
- edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
+ edge_weight = [edge_weight; ones_like(edge_weight, g.num_nodes)]
  @assert length(edge_weight) == g.num_edges
  end
  end
@@ -644,7 +643,7 @@ function tag_conv(l, g::GNNGraph, x::AbstractMatrix{T},
  if l.add_self_loops
  g = add_self_loops(g)
  if edge_weight !== nothing
- edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
+ edge_weight = [edge_weight; ones_like(edge_weight, g.num_nodes)]
  @assert length(edge_weight) == g.num_edges
  end
  end

src/layers/conv.jl

@@ -486,7 +486,7 @@ where ``\mathbf{h}^{(l)}_i`` denotes the ``l``-th hidden variables passing throu
 # Arguments
 
 - `out`: The dimension of output features.
-- `num_layers`: The number of gated recurrent unit.
+- `num_layers`: The number of recursion steps.
 - `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
 - `init`: Weight initialization function.
 
@@ -510,25 +510,25 @@ y = l(g, x)
 struct GatedGraphConv{W <: AbstractArray{<:Number, 3}, R, A} <: GNNLayer
  weight::W
  gru::R
- out_ch::Int
+ dims::Int
  num_layers::Int
  aggr::A
 end
 
 @functor GatedGraphConv
 
-function GatedGraphConv(out_ch::Int, num_layers::Int;
+function GatedGraphConv(dims::Int, num_layers::Int;
  aggr = +, init = glorot_uniform)
- w = init(out_ch, out_ch, num_layers)
- gru = GRUCell(out_ch, out_ch)
- GatedGraphConv(w, gru, out_ch, num_layers, aggr)
+ w = init(dims, dims, num_layers)
+ gru = GRUCell(dims => dims)
+ GatedGraphConv(w, gru, dims, num_layers, aggr)
 end
 
 
 (l::GatedGraphConv)(g, H) = GNNlib.gated_graph_conv(l, g, H)
 
 function Base.show(io::IO, l::GatedGraphConv)
- print(io, "GatedGraphConv(($(l.out_ch) => $(l.out_ch))^$(l.num_layers)")
+ print(io, "GatedGraphConv($(l.dims), $(l.num_layers)")
  print(io, ", aggr=", l.aggr)
  print(io, ")")
 end
@@ -1201,7 +1201,7 @@ function SGConv(ch::Pair{Int, Int}, k = 1;
  in, out = ch
  W = init(out, in)
  b = bias ? Flux.create_bias(W, true, out) : false
- SGConv(W, b, k, add_self_loops, use_edge_weight)
+ return SGConv(W, b, k, add_self_loops, use_edge_weight)
 end
 
 (l::SGConv)(g, x, edge_weight = nothing) = GNNlib.sg_conv(l, g, x, edge_weight)