skip Mooncake for some tests

Author: AstitvaAggarwal

lib/LuxLib/test/others/bmm/autodiff_tests.jl

@@ -1,7 +1,7 @@
 using LuxLib, Test
 using LuxLib: batched_matmul, batched_vec
 using NNlib: batched_adjoint, batched_transpose
-using LuxTestUtils: @test_gradients, AutoEnzyme
+using LuxTestUtils: @test_gradients, AutoEnzyme, AutoMooncake
 
 include("bmm_testsetup.jl")
 
@@ -30,15 +30,15 @@ include("bmm_testsetup.jl")
                 aType(randn(rng, Float32, P, Q, B));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
             @test_gradients(
                 fn,
                 aType(randn(rng, Float32, M, P, B)),
                 batched_transpose(aType(randn(rng, Float32, Q, P, B)));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
         end
 
@@ -57,15 +57,15 @@ include("bmm_testsetup.jl")
                 aType(randn(rng, Float32, P, Q, B));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
             @test_gradients(
                 fn,
                 aType(randn(rng, Float32, M, P)),
                 batched_adjoint(aType(randn(rng, Float32, Q, P, B)));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
 
             @test_gradients(
@@ -82,15 +82,15 @@ include("bmm_testsetup.jl")
                 aType(randn(rng, Float32, P, Q, B));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
             @test_gradients(
                 fn,
                 aType(randn(rng, Float32, M, P)),
                 batched_adjoint(aType(randn(rng, Float32, Q, P, B)));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
         end
 
@@ -109,15 +109,15 @@ include("bmm_testsetup.jl")
                 aType(randn(rng, Float32, P, Q, B));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
             @test_gradients(
                 fn,
                 aType(randn(rng, Float32, M, P, 1)),
                 batched_transpose(aType(randn(rng, Float32, Q, P, B)));
                 atol=1.0e-3,
                 rtol=1.0e-3,
-                skip_backends=[AutoEnzyme()]
+                skip_backends=[AutoEnzyme(), AutoMooncake()]
             )
         end
     end

test/misc/helpers/loss_tests.jl

@@ -156,12 +156,16 @@ end
             @jet celoss(ŷ, y)
             @jet celoss_smooth(ŷ, y)
 
+            # ŷ is Adjoint-wrapped; Mooncake can't rebuild that wrapper yet, so the
+            # gradient comes back with the wrong shape.
             @test_gradients(
                 Base.Fix2(celoss, y),
                 ŷ;
                 atol=1.0f-3,
                 rtol=1.0f-3,
-                skip_backends=VERSION ≥ v"1.11-" ? [AutoEnzyme()] : []
+                skip_backends=vcat(
+                    VERSION ≥ v"1.11-" ? [AutoEnzyme()] : [], [AutoMooncake()]
+                )
             )
         end
 
@@ -182,12 +186,15 @@ end
             @jet logitceloss(logŷ, y)
             @jet logitceloss_smooth(logŷ, y)
 
+            # Same Adjoint-wrapper issue as above.
             @test_gradients(
                 Base.Fix2(logitceloss, y),
                 logŷ;
                 atol=1.0f-3,
                 rtol=1.0f-3,
-                skip_backends=VERSION ≥ v"1.11-" ? [AutoEnzyme()] : []
+                skip_backends=vcat(
+                    VERSION ≥ v"1.11-" ? [AutoEnzyme()] : [], [AutoMooncake()]
+                )
             )
         end