diff --git a/GNNLux/test/layers/temporalconv.jl b/GNNLux/test/layers/temporalconv.jl index 4657b66df..f12a29c67 100644 --- a/GNNLux/test/layers/temporalconv.jl +++ b/GNNLux/test/layers/temporalconv.jl @@ -1,6 +1,6 @@ @testitem "layers/temporalconv" setup=[TestModuleLux] begin using .TestModuleLux - using LuxTestUtils: test_gradients, AutoTracker, AutoForwardDiff, AutoEnzyme + using LuxTestUtils: test_gradients, AutoTracker, AutoForwardDiff, AutoEnzyme, AutoMooncake rng = StableRNG(1234) g = rand_graph(rng, 10, 40) @@ -16,7 +16,7 @@ st = LuxCore.initialstates(rng, l) y1, _ = l(g, x, ps, st) 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=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) # Test with custom activation (relu) l_relu = TGCN(3=>3, act = relu) @@ -28,7 +28,7 @@ @test !isapprox(y1, y2, rtol=1.0f-2) loss_relu = (x, ps) -> sum(first(l_relu(g, x, ps, st_relu))) - test_gradients(loss_relu, x, ps_relu; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss_relu, x, ps_relu; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end @testset "A3TGCN" begin @@ -36,7 +36,7 @@ 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=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end @testset "GConvGRU" begin @@ -44,7 +44,7 @@ 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=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end @testset "GConvLSTM" begin @@ -52,7 +52,7 @@ 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=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end @testset "DCGRU" begin @@ -60,7 +60,7 @@ 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=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end @testset "EvolveGCNO" begin @@ -68,6 +68,6 @@ 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=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, tx, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end -end \ No newline at end of file +end diff --git a/GNNLux/test/test_module.jl b/GNNLux/test/test_module.jl index 3f895b996..72b867ffd 100644 --- a/GNNLux/test/test_module.jl +++ b/GNNLux/test/test_module.jl @@ -32,7 +32,7 @@ using Reexport: @reexport @reexport using StableRNGs @reexport using Random, Statistics -using LuxTestUtils: test_gradients, AutoTracker, AutoForwardDiff, AutoEnzyme +using LuxTestUtils: test_gradients, AutoTracker, AutoForwardDiff, AutoEnzyme, AutoMooncake export test_lux_layer @@ -71,7 +71,7 @@ function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x; else loss = (x, ps) -> mean(first(l(g, x, ps, st))) end - test_gradients(loss, x, ps; atol, rtol, skip_backends=[AutoForwardDiff(), AutoEnzyme()]) + test_gradients(loss, x, ps; atol, rtol, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()]) end end diff --git a/GraphNeuralNetworks/test/test_module.jl b/GraphNeuralNetworks/test/test_module.jl index 960e606e0..b4612e04f 100644 --- a/GraphNeuralNetworks/test/test_module.jl +++ b/GraphNeuralNetworks/test/test_module.jl @@ -122,7 +122,7 @@ function test_gradients( check_equal_leaves(g, g_fd; rtol, atol) end - if test_mooncake + if test_mooncake && !(graph.graph isa AbstractSparseMatrix) # Mooncake friendly tangents currently error on sparse graph internals # Mooncake gradient with respect to input via Flux integration, compared against Zygote. loss_mc_x = (xs...) -> loss(f, graph, xs...) y_mc, g_mc = Flux.withgradient(loss_mc_x, Flux.AutoMooncake(), xs...) @@ -153,7 +153,7 @@ function test_gradients( check_equal_leaves(g, g_fd; rtol, atol) end - if test_mooncake + if test_mooncake && !(graph.graph isa AbstractSparseMatrix) # Mooncake friendly tangents currently error on sparse graph internals # Mooncake gradient with respect to f via Flux integration, compared against Zygote. y_mc, g_mc = Flux.withgradient(f -> loss(f, graph, xs...), Flux.AutoMooncake(), f) @assert isapprox(y, y_mc; rtol, atol)