Add GConvGRU, GConvLSTM and DCGRU tests

JuliaGraphs · Aug 28, 2024 · a1bfcf6 · a1bfcf6
1 parent 1ee29d9
commit a1bfcf6
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/GNNLux/test/layers/temporalconv_test.jl b/GNNLux/test/layers/temporalconv_test.jl
@@ -20,4 +20,28 @@
         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 "GConvGRU" begin
+        l = GConvGRU(3=>3, 2)
+        ps = LuxCore.initialparameters(rng, l)
+        st = LuxCore.initialstates(rng, l)
+        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 "GConvLSTM" begin
+        l = GConvLSTM(3=>3, 2)
+        ps = LuxCore.initialparameters(rng, l)
+        st = LuxCore.initialstates(rng, l)
+        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 "DCGRU" begin
+        l = DCGRU(3=>3, 2)
+        ps = LuxCore.initialparameters(rng, l)
+        st = LuxCore.initialstates(rng, l)
+        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
 end