Locate tests better

Author: willtebbutt

test/ad.jl

@@ -1,4 +1,4 @@
-@testset "AD: ForwardDiff and ReverseDiff" begin
+@testset "AD: ForwardDiff, ReverseDiff, and Mooncake" begin
     @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
         f = DynamicPPL.LogDensityFunction(m)
         rand_param_values = DynamicPPL.TestUtils.rand_prior_true(m)
@@ -17,11 +17,21 @@
             θ = convert(Vector{Float64}, varinfo[:])
             logp, ref_grad = LogDensityProblems.logdensity_and_gradient(ad_forwarddiff_f, θ)
 
-            @testset "ReverseDiff with compile=$compile" for compile in (false, true)
-                adtype = ADTypes.AutoReverseDiff(; compile=compile)
-                ad_f = LogDensityProblemsAD.ADgradient(adtype, f)
-                _, grad = LogDensityProblems.logdensity_and_gradient(ad_f, θ)
-                @test grad ≈ ref_grad
+            @testset "$adtype" for adtype in [
+                ADTypes.AutoReverseDiff(; compile=false),
+                ADTypes.AutoReverseDiff(; compile=true),
+                ADTypes.AutoMooncake(; config=nothing),
+            ]
+                # Mooncake can't currently handle something that is going on in
+                # SimpleVarInfo{<:VarNamedVector}. Disable tests for now.
+                if adtype isa ADTypes.AutoMooncake &&
+                    varinfo isa DynamicPPL.SimpleVarInfo{<:DynamicPPL.VarNamedVector}
+                    @test_broken 1 == 0
+                else
+                    ad_f = LogDensityProblemsAD.ADgradient(adtype, f)
+                    _, grad = LogDensityProblems.logdensity_and_gradient(ad_f, θ)
+                    @test grad ≈ ref_grad
+                end
             end
         end
     end

test/logdensityfunction.jl

@@ -31,18 +31,6 @@ end
             θ = varinfo[:]
             @test LogDensityProblems.logdensity(logdensity, θ) ≈ logjoint(model, varinfo)
             @test LogDensityProblems.dimension(logdensity) == length(θ)
-
-            # Test a single backend on the generic
-            # ADgradient(::AbstractADType, ::LogDensityFunction) method. This really just
-            # checks that it runs at all.
-            if varinfo isa DynamicPPL.TypedVarInfo
-                ad = ADTypes.AutoMooncake(; config=nothing)
-                ∇ℓ = LogDensityProblemsAD.ADgradient(ad, logdensity)
-                @test isa(
-                    LogDensityProblems.logdensity_and_gradient(∇ℓ, θ),
-                    Tuple{Float64, Vector{Float64}},
-                )
-            end
         end
     end
 end