Skip Mooncake tests on GNNLux and Sparse graphs (#665)

* Test: Skipped Mooncake for failing layers

Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>

* skipped mooncake on sparse graphs

Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>

---------

Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>

Author: Parvm1102

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,46 +28,46 @@
         @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
         l = A3TGCN(3=>3)
         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
         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=[AutoForwardDiff(), AutoEnzyme()])
+        test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()])
     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=[AutoForwardDiff(), AutoEnzyme()])
+        test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()])
     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=[AutoForwardDiff(), AutoEnzyme()])
+        test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()])
     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=[AutoForwardDiff(), AutoEnzyme()])
+        test_gradients(loss, tx, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoForwardDiff(), AutoEnzyme(), AutoMooncake()])
     end
-end
+end

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

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)