Move ADTypeCheckContexts and optimisation tests to ad.jl as well

Author: penelopeysm

test/ad.jl

@@ -4,9 +4,11 @@ using Turing
 using DynamicPPL
 using DynamicPPL.TestUtils: DEMO_MODELS
 using DynamicPPL.TestUtils.AD: run_ad
+using Random: Random
 using StableRNGs: StableRNG
 using Test
 using ..Models: gdemo_default
+import ForwardDiff, ReverseDiff, Mooncake
 
 """Element types that are always valid for a VarInfo regardless of ADType."""
 const always_valid_eltypes = (AbstractFloat, AbstractIrrational, Integer, Rational)
@@ -181,17 +183,49 @@ ADTYPES = [
     Turing.AutoMooncake(; config=nothing),
 ]
 
+# Check that ADTypeCheckContext itself works as expected.
+@testset "ADTypeCheckContext" begin
+    @model test_model() = x ~ Normal(0, 1)
+    tm = test_model()
+    adtypes = (
+        Turing.AutoForwardDiff(),
+        Turing.AutoReverseDiff(),
+        # TODO: Mooncake
+        # Turing.AutoMooncake(config=nothing),
+    )
+    for actual_adtype in adtypes
+        sampler = Turing.HMC(0.1, 5; adtype=actual_adtype)
+        for expected_adtype in adtypes
+            contextualised_tm = DynamicPPL.contextualize(
+                tm, ADTypeCheckContext(expected_adtype, tm.context)
+            )
+            @testset "Expected: $expected_adtype, Actual: $actual_adtype" begin
+                if actual_adtype == expected_adtype
+                    # Check that this does not throw an error.
+                    Turing.sample(contextualised_tm, sampler, 2)
+                else
+                    @test_throws AbstractWrongADBackendError Turing.sample(
+                        contextualised_tm, sampler, 2
+                    )
+                end
+            end
+        end
+    end
+end
+
 @testset verbose = true "AD / ADTypeCheckContext" begin
-    # This testset ensures that samplers don't accidentally override the AD
-    # backend set in it.
-    @testset "Check ADType" begin
+    # This testset ensures that samplers or optimisers don't accidentally
+    # override the AD backend set in it.
+    @testset "adtype=$adtype" for adtype in ADTYPES
         seed = 123
         alg = HMC(0.1, 10; adtype=adtype)
         m = DynamicPPL.contextualize(
             gdemo_default, ADTypeCheckContext(adtype, gdemo_default.context)
         )
         # These will error if the adbackend being used is not the one set.
         sample(StableRNG(seed), m, alg, 10)
+        maximum_likelihood(m; adtype=adtype)
+        maximum_a_posteriori(m; adtype=adtype)
     end
 end
 

test/mcmc/Inference.jl

@@ -2,7 +2,6 @@ module InferenceTests
 
 using ..Models: gdemo_d, gdemo_default
 using ..NumericalTests: check_gdemo, check_numerical
-import ..ADUtils
 using Distributions: Bernoulli, Beta, InverseGamma, Normal
 using Distributions: sample
 import DynamicPPL

test/mcmc/abstractmcmc.jl

@@ -1,6 +1,5 @@
 module AbstractMCMCTests
 
-import ..ADUtils
 using AbstractMCMC: AbstractMCMC
 using AdvancedMH: AdvancedMH
 using Distributions: sample

test/mcmc/gibbs.jl

@@ -7,7 +7,6 @@ using ..NumericalTests:
     check_gdemo,
     check_numerical,
     two_sample_test
-import ..ADUtils
 import Combinatorics
 using AbstractMCMC: AbstractMCMC
 using Distributions: InverseGamma, Normal

test/mcmc/hmc.jl

@@ -1,9 +1,7 @@
 module HMCTests
 
 using ..Models: gdemo_default
-using ..ADUtils: ADTypeCheckContext
 using ..NumericalTests: check_gdemo, check_numerical
-import ..ADUtils
 using Bijectors: Bijectors
 using Distributions: Bernoulli, Beta, Categorical, Dirichlet, Normal, Wishart, sample
 using DynamicPPL: DynamicPPL, Sampler

test/mcmc/sghmc.jl

@@ -2,7 +2,6 @@ module SGHMCTests
 
 using ..Models: gdemo_default
 using ..NumericalTests: check_gdemo
-import ..ADUtils
 using DynamicPPL.TestUtils.AD: run_ad
 using DynamicPPL.TestUtils: DEMO_MODELS
 using DynamicPPL: DynamicPPL

test/optimisation/Optimisation.jl

@@ -1,7 +1,6 @@
 module OptimisationTests
 
 using ..Models: gdemo, gdemo_default
-using ..ADUtils: ADUtils
 using Distributions
 using Distributions.FillArrays: Zeros
 using DynamicPPL: DynamicPPL
@@ -624,16 +623,6 @@ using Turing
         @assert get(result, :c) == (; :c => Array{Float64}[])
     end
 
-    @testset "ADType test with $adbackend" for adbackend in ADUtils.adbackends
-        Random.seed!(222)
-        m = DynamicPPL.contextualize(
-            gdemo_default, ADUtils.ADTypeCheckContext(adbackend, gdemo_default.context)
-        )
-        # These will error if the adbackend being used is not the one set.
-        maximum_likelihood(m; adtype=adbackend)
-        maximum_a_posteriori(m; adtype=adbackend)
-    end
-
     @testset "Collinear coeftable" begin
         xs = [-1.0, 0.0, 1.0]
         ys = [0.0, 0.0, 0.0]

test/test_utils/test_utils.jl

@@ -8,34 +8,4 @@ using Test: @test, @testset, @test_throws
 using Turing: Turing
 using Turing: DynamicPPL
 
-# Check that the ADTypeCheckContext works as expected.
-@testset "ADTypeCheckContext" begin
-    Turing.@model test_model() = x ~ Turing.Normal(0, 1)
-    tm = test_model()
-    adtypes = (
-        Turing.AutoForwardDiff(),
-        Turing.AutoReverseDiff(),
-        # TODO: Mooncake
-        # Turing.AutoMooncake(config=nothing),
-    )
-    for actual_adtype in adtypes
-        sampler = Turing.HMC(0.1, 5; adtype=actual_adtype)
-        for expected_adtype in adtypes
-            contextualised_tm = DynamicPPL.contextualize(
-                tm, ADTypeCheckContext(expected_adtype, tm.context)
-            )
-            @testset "Expected: $expected_adtype, Actual: $actual_adtype" begin
-                if actual_adtype == expected_adtype
-                    # Check that this does not throw an error.
-                    Turing.sample(contextualised_tm, sampler, 2)
-                else
-                    @test_throws AbstractWrongADBackendError Turing.sample(
-                        contextualised_tm, sampler, 2
-                    )
-                end
-            end
-        end
-    end
-end
-
 end