Add Flux.Recur and TGCN (#319)

* Add `TGCN` and `Flux.Recur`

* Add docstring `TGCN`

* Add comment

* Specify type

* Add test

* FIx identation

* Add spaces

Co-authored-by: Carlo Lucibello <[email protected]>

* Improved docstring

Co-authored-by: Carlo Lucibello <[email protected]>

* Fix docstring

Co-authored-by: Carlo Lucibello <[email protected]>

* Add `_applylayer`

---------

Co-authored-by: Carlo Lucibello <[email protected]>

Author: aurorarossi

src/GraphNeuralNetworks.jl

@@ -71,7 +71,7 @@ export
       HeteroGraphConv,
 
 # layers/temporalconv
-      TGCNCell,
+      TGCN,
 
 # layers/pool
       GlobalPool,

src/layers/temporalconv.jl

@@ -1,22 +1,15 @@
-"""
-    TGCNCell(in => out; [bias, init, add_self_loops, use_edge_weight])
-
-Temporal Graph Convolutional Network (T-GCN) cell from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf).
-
-Performs a layer of GCNConv to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies.
-
-# Arguments
+# Adapting Flux.Recur to work with GNNGraphs
+function (m::Flux.Recur)(g::GNNGraph, x)
+    m.state, y = m.cell(m.state, g, x)
+    return y
+end
+    
+function (m::Flux.Recur)(g::GNNGraph, x::AbstractArray{T, 3}) where T
+    h = [m(g, x_t) for x_t in Flux.eachlastdim(x)]
+    sze = size(h[1])
+    reshape(reduce(hcat, h), sze[1], sze[2], length(h))
+end
 
-- `in`: Number of input features.
-- `out`: Number of output features.
-- `bias`: Add learnable bias. Default `true`.
-- `init`: Weights' initializer. Default `glorot_uniform`.
-- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
-- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available).
-                     If `add_self_loops=true` the new weights will be set to 1. 
-                     This option is ignored if the `edge_weight` is explicitly provided in the forward pass.
-                     Default `false`.
-"""
 struct TGCNCell <: GNNLayer
     conv::GCNConv
     gru::Flux.GRUv3Cell
@@ -50,3 +43,61 @@ end
 function Base.show(io::IO, tgcn::TGCNCell)
     print(io, "TGCNCell($(tgcn.in) => $(tgcn.out))")
 end
+
+"""
+    TGCN(in => out; [bias, init, add_self_loops, use_edge_weight])
+
+Temporal Graph Convolutional Network (T-GCN) recurrent layer from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf).
+
+Performs a layer of GCNConv to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies.
+
+# Arguments
+
+- `in`: Number of input features.
+- `out`: Number of output features.
+- `bias`: Add learnable bias. Default `true`.
+- `init`: Weights' initializer. Default `glorot_uniform`.
+- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
+- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available).
+                     If `add_self_loops=true` the new weights will be set to 1. 
+                     This option is ignored if the `edge_weight` is explicitly provided in the forward pass.
+                     Default `false`.
+# Examples
+
+```jldoctest
+julia> tgcn = TGCN(2 => 6)
+Recur(
+  TGCNCell(
+    GCNConv(2 => 6, σ),                 # 18 parameters
+    GRUv3Cell(6 => 6),                  # 240 parameters
+    Float32[0.0; 0.0; … ; 0.0; 0.0;;],  # 6 parameters  (all zero)
+    2,
+    6,
+  ),
+)         # Total: 8 trainable arrays, 264 parameters,
+          # plus 1 non-trainable, 6 parameters, summarysize 1.492 KiB.
+
+julia> g, x = rand_graph(5, 10), rand(Float32, 2, 5);
+
+julia> y = tgcn(g, x);
+
+julia> size(y)
+(6, 5)
+
+julia> Flux.reset!(tgcn);
+
+julia> tgcn(rand_graph(5, 10), rand(Float32, 2, 5, 20)) |> size # batch size of 20
+(6, 5, 20)
+```
+
+!!! warning "Batch size changes"
+    Failing to call `reset!` when the input batch size changes can lead to unexpected behavior.
+"""
+TGCN(ch; kwargs...) = Flux.Recur(TGCNCell(ch; kwargs...))
+
+Flux.Recur(tgcn::TGCNCell) = Flux.Recur(tgcn, tgcn.state0)
+
+# make TGCN compatible with GNNChain
+(l::Flux.Recur{TGCNCell})(g::GNNGraph) = GNNGraph(g, ndata = l(g, node_features(g)))
+_applylayer(l::Flux.Recur{TGCNCell}, g::GNNGraph, x) = l(g, x)
+_applylayer(l::Flux.Recur{TGCNCell}, g::GNNGraph) = l(g)

test/layers/temporalconv.jl

@@ -8,9 +8,17 @@ g1 = GNNGraph(rand_graph(N,8),
                 graph_type = :sparse)
 
 @testset "TGCNCell" begin
-    tgcn = TGCNCell(in_channel => out_channel)
+    tgcn = GraphNeuralNetworks.TGCNCell(in_channel => out_channel)
     h, x̃ = tgcn(tgcn.state0, g1, g1.ndata.x)
     @test size(h) == (out_channel, N)
     @test size(x̃) == (out_channel, N)
     @test h == x̃
+end
+
+@testset "TGCN" begin
+    tgcn = TGCN(in_channel => out_channel)
+    @test size(Flux.gradient(x -> sum(tgcn(g1, x)), g1.ndata.x)[1]) == (in_channel, N)
+    model = GNNChain(TGCN(in_channel => out_channel), Dense(out_channel, 1))
+    @test size(model(g1, g1.ndata.x)) == (1, N)
+    @test model(g1) isa GNNGraph            
 end