Adapt GINConv to TemporalSnapshotsGNNGraphs (#376)

* Adapt GINConv to TemporalSnapshotsGNNGraphs

* Add test GINConv

Author: aurorarossi

src/layers/temporalconv.jl

@@ -186,3 +186,7 @@ end
 function Base.show(io::IO, a3tgcn::A3TGCN)
     print(io, "A3TGCN($(a3tgcn.in) => $(a3tgcn.out))")
 end
+
+function (l::GINConv)(tg::TemporalSnapshotsGNNGraph, x::AbstractVector)
+    return l.(tg.snapshots, x)
+end

test/layers/temporalconv.jl

@@ -1,12 +1,15 @@
 in_channel = 3
 out_channel = 5
 N = 4
+S = 5
 T = Float32
 
 g1 = GNNGraph(rand_graph(N,8),
                 ndata = rand(T, in_channel, N),
                 graph_type = :sparse)
 
+tg = TemporalSnapshotsGNNGraph([g1 for _ in 1:S])
+
 @testset "TGCNCell" begin
     tgcn = GraphNeuralNetworks.TGCNCell(in_channel => out_channel)
     h, x̃ = tgcn(tgcn.state0, g1, g1.ndata.x)
@@ -29,4 +32,11 @@ end
     model = GNNChain(A3TGCN(in_channel => out_channel), Dense(out_channel, 1))
     @test size(model(g1, g1.ndata.x)) == (1, N)
     @test model(g1) isa GNNGraph            
+end
+
+@testset "GINConv" begin
+    ginconv = GINConv(Dense(in_channel => out_channel),0.3)
+    @test length(ginconv(tg, tg.ndata.x)) == S
+    @test size(ginconv(tg, tg.ndata.x)[1]) == (out_channel, N) 
+    @test length(Flux.gradient(x ->sum(sum(ginconv(tg, x))), tg.ndata.x)[1]) == S    
 end