Add EvolveGCNO test

JuliaGraphs · Sep 6, 2024 · d49bde7 · d49bde7
1 parent 1e451c4
commit d49bde7
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/GNNLux/test/layers/temporalconv_test.jl b/GNNLux/test/layers/temporalconv_test.jl
@@ -5,6 +5,9 @@
     g = rand_graph(rng, 10, 40)
     x = randn(rng, Float32, 3, 10)
 
+    tg = TemporalSnapshotsGNNGraph([g for _ in 1:5])
+    tx = [x for _ in 1:5]
+
     @testset "TGCN" begin
         l = TGCN(3=>3)
         ps = LuxCore.initialparameters(rng, l)
@@ -44,4 +47,12 @@
         loss = (x, ps) -> sum(first(l(g, x, ps, st)))
         test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
     end
+
+    @testset "EvolveGCNO" begin
+        l = EvolveGCNO(3=>3)
+        ps = LuxCore.initialparameters(rng, l)
+        st = LuxCore.initialstates(rng, l)
+        loss = (tx, ps) -> sum(sum(first(l(tg, tx, ps, st))))
+        test_gradients(loss, tx, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
+    end
 end