From bb1bcefe68d2708fef90566062e67ba2968ceab6 Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 14:28:35 +0100
Subject: [PATCH 01/19] Test with Tapir

---
 test/Project.toml              | 2 ++
 test/mcmc/Inference.jl         | 4 +++-
 test/mcmc/abstractmcmc.jl      | 5 +++--
 test/mcmc/gibbs.jl             | 6 +++---
 test/mcmc/gibbs_conditional.jl | 6 +++---
 test/mcmc/hmc.jl               | 4 +++-
 test/mcmc/sghmc.jl             | 6 ++++--
 test/runtests.jl               | 1 +
 8 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/test/Project.toml b/test/Project.toml
index 67292d2af..01359514c 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -30,6 +30,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+Tapir = "07d77754-e150-4737-8c94-cd238a1fb45b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
@@ -67,6 +68,7 @@ SpecialFunctions = "0.10.3, 1, 2"
 StableRNGs = "1"
 StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
+Tapir = "0.2.24"
 TimerOutputs = "0.5"
 Tracker = "0.2.11"
 Zygote = "0.5.4, 0.6"
diff --git a/test/mcmc/Inference.jl b/test/mcmc/Inference.jl
index 64a5e95df..8350339d4 100644
--- a/test/mcmc/Inference.jl
+++ b/test/mcmc/Inference.jl
@@ -2,6 +2,7 @@ module InferenceTests
 
 using ..Models: gdemo_d, gdemo_default
 using ..NumericalTests: check_gdemo, check_numerical
+using ..ADUtils: adbackends
 using Distributions: Bernoulli, Beta, InverseGamma, Normal
 using Distributions: sample
 import DynamicPPL
@@ -11,10 +12,11 @@ using LinearAlgebra: I
 import MCMCChains
 import Random
 import ReverseDiff
+import Tapir
 using Test: @test, @test_throws, @testset
 using Turing
 
-@testset "Testing inference.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing inference.jl with $adbackend" for adbackend in adbackends
     # Only test threading if 1.3+.
     if VERSION > v"1.2"
         @testset "threaded sampling" begin
diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index 10d202da3..20cfe4b63 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -1,5 +1,6 @@
 module AbstractMCMCTests
 
+using ..ADUtils: adbackends
 using AdvancedMH: AdvancedMH
 using Distributions: sample
 using Distributions.FillArrays: Zeros
@@ -11,6 +12,7 @@ using LogDensityProblemsAD: LogDensityProblemsAD
 using Random: Random
 using ReverseDiff: ReverseDiff
 using StableRNGs: StableRNG
+import Tapir
 using Test: @test, @test_throws, @testset
 using Turing
 using Turing.Inference: AdvancedHMC
@@ -112,8 +114,7 @@ end
 
 @testset "External samplers" begin
     @testset "AdvancedHMC.jl" begin
-        # Try a few different AD backends.
-        @testset "adtype=$adtype" for adtype in [AutoForwardDiff(), AutoReverseDiff()]
+        @testset "adtype=$adtype" for adtype in adbackends
             @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
                 # Need some functionality to initialize the sampler.
                 # TODO: Remove this once the constructors in the respective packages become "lazy".
diff --git a/test/mcmc/gibbs.jl b/test/mcmc/gibbs.jl
index f30dc0f77..818efedac 100644
--- a/test/mcmc/gibbs.jl
+++ b/test/mcmc/gibbs.jl
@@ -2,19 +2,19 @@ module GibbsTests
 
 using ..Models: MoGtest_default, gdemo, gdemo_default
 using ..NumericalTests: check_MoGtest_default, check_gdemo, check_numerical
+using ..ADUtils: adbackends
 using Distributions: InverseGamma, Normal
 using Distributions: sample
 using ForwardDiff: ForwardDiff
 using Random: Random
 using ReverseDiff: ReverseDiff
+import Tapir
 using Test: @test, @testset
 using Turing
 using Turing: Inference
 using Turing.RandomMeasures: ChineseRestaurantProcess, DirichletProcess
 
-@testset "Testing gibbs.jl with $adbackend" for adbackend in (
-    AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false)
-)
+@testset "Testing gibbs.jl with $adbackend" for adbackend in adbackends
     @testset "gibbs constructor" begin
         N = 500
         s1 = Gibbs(HMC(0.1, 5, :s, :m; adtype=adbackend))
diff --git a/test/mcmc/gibbs_conditional.jl b/test/mcmc/gibbs_conditional.jl
index 6110bbdb1..11b4ff78b 100644
--- a/test/mcmc/gibbs_conditional.jl
+++ b/test/mcmc/gibbs_conditional.jl
@@ -2,6 +2,7 @@ module GibbsConditionalTests
 
 using ..Models: gdemo, gdemo_default
 using ..NumericalTests: check_gdemo, check_numerical
+using ..ADUtils: adbackends
 using Clustering: Clustering
 using Distributions: Categorical, InverseGamma, Normal, sample
 using ForwardDiff: ForwardDiff
@@ -11,12 +12,11 @@ using ReverseDiff: ReverseDiff
 using StableRNGs: StableRNG
 using StatsBase: counts
 using StatsFuns: StatsFuns
+import Tapir
 using Test: @test, @testset
 using Turing
 
-@testset "Testing gibbs conditionals.jl with $adbackend" for adbackend in (
-    AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false)
-)
+@testset "Testing gibbs conditionals.jl with $adbackend" for adbackend in adbackends
     Random.seed!(1000)
     rng = StableRNG(123)
 
diff --git a/test/mcmc/hmc.jl b/test/mcmc/hmc.jl
index 968f24d7b..283ec2264 100644
--- a/test/mcmc/hmc.jl
+++ b/test/mcmc/hmc.jl
@@ -3,6 +3,7 @@ module HMCTests
 using ..Models: gdemo_default
 #using ..Models: gdemo
 using ..NumericalTests: check_gdemo, check_numerical
+using ..ADUtils: adbackends
 using Distributions: Bernoulli, Beta, Categorical, Dirichlet, Normal, Wishart, sample
 import DynamicPPL
 using DynamicPPL: Sampler
@@ -13,10 +14,11 @@ using LinearAlgebra: I, dot, vec
 import Random
 using StableRNGs: StableRNG
 using StatsFuns: logistic
+import Tapir
 using Test: @test, @test_logs, @testset
 using Turing
 
-@testset "Testing hmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing hmc.jl with $adbackend" for adbackend in adbackend
     # Set a seed
     rng = StableRNG(123)
     @testset "constrained bounded" begin
diff --git a/test/mcmc/sghmc.jl b/test/mcmc/sghmc.jl
index a5829eb18..104f171c6 100644
--- a/test/mcmc/sghmc.jl
+++ b/test/mcmc/sghmc.jl
@@ -2,15 +2,17 @@ module SGHMCTests
 
 using ..Models: gdemo_default
 using ..NumericalTests: check_gdemo
+using ..ADUtils: adbackends
 using Distributions: sample
 import ForwardDiff
 using LinearAlgebra: dot
 import ReverseDiff
 using StableRNGs: StableRNG
+import Tapir
 using Test: @test, @testset
 using Turing
 
-@testset "Testing sghmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing sghmc.jl with $adbackend" for adbackend in adbackends
     @testset "sghmc constructor" begin
         alg = SGHMC(; learning_rate=0.01, momentum_decay=0.1, adtype=adbackend)
         @test alg isa SGHMC
@@ -36,7 +38,7 @@ using Turing
     end
 end
 
-@testset "Testing sgld.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing sgld.jl with $adbackend" for adbackend in adbackends
     @testset "sgld constructor" begin
         alg = SGLD(; stepsize=PolynomialStepsize(0.25), adtype=adbackend)
         @test alg isa SGLD
diff --git a/test/runtests.jl b/test/runtests.jl
index 48a00122d..1aa8bb635 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,6 +7,7 @@ import Turing
 
 include(pkgdir(Turing) * "/test/test_utils/models.jl")
 include(pkgdir(Turing) * "/test/test_utils/numerical_tests.jl")
+include(pkgdir(Turing) * "/test/test_utils/ad_utils.jl")
 
 Turing.setprogress!(false)
 

From 596457a801a2d70e5f0bf0a6fd8497d3f8b572cd Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 15:02:07 +0100
Subject: [PATCH 02/19] Relax Tapir version bound

---
 test/Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Project.toml b/test/Project.toml
index 01359514c..a2727f3e0 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -68,7 +68,7 @@ SpecialFunctions = "0.10.3, 1, 2"
 StableRNGs = "1"
 StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
-Tapir = "0.2.24"
+Tapir = "0.2"
 TimerOutputs = "0.5"
 Tracker = "0.2.11"
 Zygote = "0.5.4, 0.6"

From a4ef6b37ec81884430c5e0d621139a57e6aac6d7 Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 15:05:39 +0100
Subject: [PATCH 03/19] Relax Tapir version bounds more

---
 test/Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Project.toml b/test/Project.toml
index a2727f3e0..7517bc424 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -68,7 +68,7 @@ SpecialFunctions = "0.10.3, 1, 2"
 StableRNGs = "1"
 StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
-Tapir = "0.2"
+Tapir = "0.1, 0.2"
 TimerOutputs = "0.5"
 Tracker = "0.2.11"
 Zygote = "0.5.4, 0.6"

From 78af8ad3685ff2afc5edf68fa89bbe42c2946c7f Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 15:10:22 +0100
Subject: [PATCH 04/19] Add test/test_utils/ad_utils.jl

---
 test/test_utils/ad_utils.jl | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)
 create mode 100644 test/test_utils/ad_utils.jl

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
new file mode 100644
index 000000000..5d2687456
--- /dev/null
+++ b/test/test_utils/ad_utils.jl
@@ -0,0 +1,16 @@
+module ADUtils
+
+import Turing
+
+"""
+All the ADTypes on which we want to run the tests.
+"""
+adbackends = [
+    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
+]
+# AutoTapir isn't supported for older Julia versions, hence the check.
+if isdefined(Turing.AutoTapir)
+    push!(backends, Turing.AutoTapir())
+end
+
+end

From 5796b1302fe4b57876dc62d4472c4b73308a1c42 Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 15:44:00 +0100
Subject: [PATCH 05/19] Change how Tapir is installed for tests

---
 test/Project.toml              |  2 --
 test/mcmc/Inference.jl         |  7 ++++---
 test/mcmc/abstractmcmc.jl      |  7 ++++---
 test/mcmc/gibbs.jl             |  7 ++++---
 test/mcmc/gibbs_conditional.jl |  7 ++++---
 test/mcmc/hmc.jl               |  7 ++++---
 test/mcmc/sghmc.jl             |  9 +++++----
 test/test_utils/ad_utils.jl    | 11 ++++++++---
 8 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/test/Project.toml b/test/Project.toml
index 7517bc424..67292d2af 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -30,7 +30,6 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
-Tapir = "07d77754-e150-4737-8c94-cd238a1fb45b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
@@ -68,7 +67,6 @@ SpecialFunctions = "0.10.3, 1, 2"
 StableRNGs = "1"
 StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
-Tapir = "0.1, 0.2"
 TimerOutputs = "0.5"
 Tracker = "0.2.11"
 Zygote = "0.5.4, 0.6"
diff --git a/test/mcmc/Inference.jl b/test/mcmc/Inference.jl
index 8350339d4..15ec6149c 100644
--- a/test/mcmc/Inference.jl
+++ b/test/mcmc/Inference.jl
@@ -2,7 +2,7 @@ module InferenceTests
 
 using ..Models: gdemo_d, gdemo_default
 using ..NumericalTests: check_gdemo, check_numerical
-using ..ADUtils: adbackends
+import ..ADUtils
 using Distributions: Bernoulli, Beta, InverseGamma, Normal
 using Distributions: sample
 import DynamicPPL
@@ -12,11 +12,12 @@ using LinearAlgebra: I
 import MCMCChains
 import Random
 import ReverseDiff
-import Tapir
 using Test: @test, @test_throws, @testset
 using Turing
 
-@testset "Testing inference.jl with $adbackend" for adbackend in adbackends
+ADUtils.install_tapir && import Tapir
+
+@testset "Testing inference.jl with $adbackend" for adbackend in ADUtils.adbackends
     # Only test threading if 1.3+.
     if VERSION > v"1.2"
         @testset "threaded sampling" begin
diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index 20cfe4b63..0aae153dd 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -1,6 +1,6 @@
 module AbstractMCMCTests
 
-using ..ADUtils: adbackends
+import ..ADUtils
 using AdvancedMH: AdvancedMH
 using Distributions: sample
 using Distributions.FillArrays: Zeros
@@ -12,11 +12,12 @@ using LogDensityProblemsAD: LogDensityProblemsAD
 using Random: Random
 using ReverseDiff: ReverseDiff
 using StableRNGs: StableRNG
-import Tapir
 using Test: @test, @test_throws, @testset
 using Turing
 using Turing.Inference: AdvancedHMC
 
+ADUtils.install_tapir && import Tapir
+
 function initialize_nuts(model::Turing.Model)
     # Create a log-density function with an implementation of the
     # gradient so we ensure that we're using the same AD backend as in Turing.
@@ -114,7 +115,7 @@ end
 
 @testset "External samplers" begin
     @testset "AdvancedHMC.jl" begin
-        @testset "adtype=$adtype" for adtype in adbackends
+        @testset "adtype=$adtype" for adtype in ADUtils.adbackends
             @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
                 # Need some functionality to initialize the sampler.
                 # TODO: Remove this once the constructors in the respective packages become "lazy".
diff --git a/test/mcmc/gibbs.jl b/test/mcmc/gibbs.jl
index 818efedac..6868cb5e8 100644
--- a/test/mcmc/gibbs.jl
+++ b/test/mcmc/gibbs.jl
@@ -2,19 +2,20 @@ module GibbsTests
 
 using ..Models: MoGtest_default, gdemo, gdemo_default
 using ..NumericalTests: check_MoGtest_default, check_gdemo, check_numerical
-using ..ADUtils: adbackends
+import ..ADUtils
 using Distributions: InverseGamma, Normal
 using Distributions: sample
 using ForwardDiff: ForwardDiff
 using Random: Random
 using ReverseDiff: ReverseDiff
-import Tapir
 using Test: @test, @testset
 using Turing
 using Turing: Inference
 using Turing.RandomMeasures: ChineseRestaurantProcess, DirichletProcess
 
-@testset "Testing gibbs.jl with $adbackend" for adbackend in adbackends
+ADUtils.install_tapir && import Tapir
+
+@testset "Testing gibbs.jl with $adbackend" for adbackend in ADUtils.adbackends
     @testset "gibbs constructor" begin
         N = 500
         s1 = Gibbs(HMC(0.1, 5, :s, :m; adtype=adbackend))
diff --git a/test/mcmc/gibbs_conditional.jl b/test/mcmc/gibbs_conditional.jl
index 11b4ff78b..3f02c7594 100644
--- a/test/mcmc/gibbs_conditional.jl
+++ b/test/mcmc/gibbs_conditional.jl
@@ -2,7 +2,7 @@ module GibbsConditionalTests
 
 using ..Models: gdemo, gdemo_default
 using ..NumericalTests: check_gdemo, check_numerical
-using ..ADUtils: adbackends
+import ..ADUtils
 using Clustering: Clustering
 using Distributions: Categorical, InverseGamma, Normal, sample
 using ForwardDiff: ForwardDiff
@@ -12,11 +12,12 @@ using ReverseDiff: ReverseDiff
 using StableRNGs: StableRNG
 using StatsBase: counts
 using StatsFuns: StatsFuns
-import Tapir
 using Test: @test, @testset
 using Turing
 
-@testset "Testing gibbs conditionals.jl with $adbackend" for adbackend in adbackends
+ADUtils.install_tapir && import Tapir
+
+@testset "Testing gibbs conditionals.jl with $adbackend" for adbackend in ADUtils.adbackends
     Random.seed!(1000)
     rng = StableRNG(123)
 
diff --git a/test/mcmc/hmc.jl b/test/mcmc/hmc.jl
index 283ec2264..37ffd08a1 100644
--- a/test/mcmc/hmc.jl
+++ b/test/mcmc/hmc.jl
@@ -3,7 +3,7 @@ module HMCTests
 using ..Models: gdemo_default
 #using ..Models: gdemo
 using ..NumericalTests: check_gdemo, check_numerical
-using ..ADUtils: adbackends
+import ..ADUtils
 using Distributions: Bernoulli, Beta, Categorical, Dirichlet, Normal, Wishart, sample
 import DynamicPPL
 using DynamicPPL: Sampler
@@ -14,11 +14,12 @@ using LinearAlgebra: I, dot, vec
 import Random
 using StableRNGs: StableRNG
 using StatsFuns: logistic
-import Tapir
 using Test: @test, @test_logs, @testset
 using Turing
 
-@testset "Testing hmc.jl with $adbackend" for adbackend in adbackend
+ADUtils.install_tapir && import Tapir
+
+@testset "Testing hmc.jl with $adbackend" for adbackend in ADUtils.adbackends
     # Set a seed
     rng = StableRNG(123)
     @testset "constrained bounded" begin
diff --git a/test/mcmc/sghmc.jl b/test/mcmc/sghmc.jl
index 104f171c6..95b3bc543 100644
--- a/test/mcmc/sghmc.jl
+++ b/test/mcmc/sghmc.jl
@@ -2,17 +2,18 @@ module SGHMCTests
 
 using ..Models: gdemo_default
 using ..NumericalTests: check_gdemo
-using ..ADUtils: adbackends
+import ..ADUtils
 using Distributions: sample
 import ForwardDiff
 using LinearAlgebra: dot
 import ReverseDiff
 using StableRNGs: StableRNG
-import Tapir
 using Test: @test, @testset
 using Turing
 
-@testset "Testing sghmc.jl with $adbackend" for adbackend in adbackends
+ADUtils.install_tapir && import Tapir
+
+@testset "Testing sghmc.jl with $adbackend" for adbackend in ADUtils.adbackends
     @testset "sghmc constructor" begin
         alg = SGHMC(; learning_rate=0.01, momentum_decay=0.1, adtype=adbackend)
         @test alg isa SGHMC
@@ -38,7 +39,7 @@ using Turing
     end
 end
 
-@testset "Testing sgld.jl with $adbackend" for adbackend in adbackends
+@testset "Testing sgld.jl with $adbackend" for adbackend in ADUtils.adbackends
     @testset "sgld constructor" begin
         alg = SGLD(; stepsize=PolynomialStepsize(0.25), adtype=adbackend)
         @test alg isa SGLD
diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index 5d2687456..9b95cfaf4 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -1,5 +1,6 @@
 module ADUtils
 
+import Pkg
 import Turing
 
 """
@@ -8,9 +9,13 @@ All the ADTypes on which we want to run the tests.
 adbackends = [
     Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
 ]
-# AutoTapir isn't supported for older Julia versions, hence the check.
-if isdefined(Turing.AutoTapir)
-    push!(backends, Turing.AutoTapir())
+
+# Tapir isn't supported for older Julia versions, hence the check.
+install_tapir = isdefined(Turing.AutoTapir)
+if install_tapir
+    # TODO(mhauru) Is there a better way to install optional dependencies like this?
+    Pkg.add("Tapir")
+    push!(adbackends, Turing.AutoTapir())
 end
 
 end

From 253cdc1746455c9a7e6ccfa3e4001e7f859d0d0f Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 16:01:24 +0100
Subject: [PATCH 06/19] Typo fix

---
 test/test_utils/ad_utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index 9b95cfaf4..09fe21656 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -11,7 +11,7 @@ adbackends = [
 ]
 
 # Tapir isn't supported for older Julia versions, hence the check.
-install_tapir = isdefined(Turing.AutoTapir)
+install_tapir = isdefined(Turing, :AutoTapir)
 if install_tapir
     # TODO(mhauru) Is there a better way to install optional dependencies like this?
     Pkg.add("Tapir")

From e50d8c553f6fd6a9fb96e27e3eabe76d029952bd Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Wed, 17 Jul 2024 16:48:25 +0100
Subject: [PATCH 07/19] Turn Tapir's safe mode off

---
 test/test_utils/ad_utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index 09fe21656..375ea3da2 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -15,7 +15,7 @@ install_tapir = isdefined(Turing, :AutoTapir)
 if install_tapir
     # TODO(mhauru) Is there a better way to install optional dependencies like this?
     Pkg.add("Tapir")
-    push!(adbackends, Turing.AutoTapir())
+    push!(adbackends, Turing.AutoTapir(false))
 end
 
 end

From 2af71fe3b367089c44eff0dc1409af894010f393 Mon Sep 17 00:00:00 2001
From: Markus Hauru <markus@mhauru.org>
Date: Thu, 18 Jul 2024 09:28:26 +0100
Subject: [PATCH 08/19] Use standard AutoReverseDiff constructor

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>
---
 test/test_utils/ad_utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index 375ea3da2..3380b4573 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -7,7 +7,7 @@ import Turing
 All the ADTypes on which we want to run the tests.
 """
 adbackends = [
-    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
+    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff{false}()
 ]
 
 # Tapir isn't supported for older Julia versions, hence the check.

From f376f90724a14a6b7dea997e275b146fb33e967f Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Thu, 18 Jul 2024 10:05:30 +0100
Subject: [PATCH 09/19] Revert back to previous AutoReverseDiff constructor

---
 test/test_utils/ad_utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index 3380b4573..375ea3da2 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -7,7 +7,7 @@ import Turing
 All the ADTypes on which we want to run the tests.
 """
 adbackends = [
-    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff{false}()
+    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
 ]
 
 # Tapir isn't supported for older Julia versions, hence the check.

From 15fd3e3947036d8cf3d1f11d2a46898e5f444946 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 31 Jul 2024 11:34:37 +0100
Subject: [PATCH 10/19] modify `setvarinfo`

---
 src/mcmc/abstractmcmc.jl  | 14 ++++++++++----
 test/mcmc/abstractmcmc.jl |  4 +++-
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/mcmc/abstractmcmc.jl b/src/mcmc/abstractmcmc.jl
index 1373a4bbd..d1fcb4308 100644
--- a/src/mcmc/abstractmcmc.jl
+++ b/src/mcmc/abstractmcmc.jl
@@ -40,11 +40,17 @@ getstats(transition::AdvancedHMC.Transition) = transition.stat
 getparams(::DynamicPPL.Model, transition::AdvancedMH.Transition) = transition.params
 
 getvarinfo(f::DynamicPPL.LogDensityFunction) = f.varinfo
-getvarinfo(f::LogDensityProblemsAD.ADGradientWrapper) = getvarinfo(parent(f))
+function getvarinfo(f::LogDensityProblemsAD.ADGradientWrapper)
+    return getvarinfo(LogDensityProblemsAD.parent(f))
+end
 
 setvarinfo(f::DynamicPPL.LogDensityFunction, varinfo) = Accessors.@set f.varinfo = varinfo
-function setvarinfo(f::LogDensityProblemsAD.ADGradientWrapper, varinfo)
-    return Accessors.@set f.ℓ = setvarinfo(f.ℓ, varinfo)
+function setvarinfo(
+    f::LogDensityProblemsAD.ADGradientWrapper, varinfo, adtype::ADTypes.AbstractADType
+)
+    return LogDensityProblemsAD.ADgradient(
+        adtype, setvarinfo(LogDensityProblemsAD.parent(f), varinfo)
+    )
 end
 
 """
@@ -120,7 +126,7 @@ function AbstractMCMC.step(
             varinfo = DynamicPPL.link(varinfo, model)
         end
     end
-    f = setvarinfo(f, varinfo)
+    f = setvarinfo(f, varinfo, sampler.alg.adtype)
 
     # Then just call `AdvancedHMC.step` with the right arguments.
     if initial_state === nothing
diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index 0aae153dd..e35df1260 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -25,7 +25,9 @@ function initialize_nuts(model::Turing.Model)
 
     # Link the varinfo.
     f = Turing.Inference.setvarinfo(
-        f, DynamicPPL.link!!(Turing.Inference.getvarinfo(f), model)
+        f,
+        DynamicPPL.link!!(Turing.Inference.getvarinfo(f), model),
+        getADType(DynamicPPL.getcontext(f)),
     )
 
     # Choose parameter dimensionality and initial parameter value

From 311e497548a9c56047345835a47e1dc5cf5145c6 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 31 Jul 2024 13:43:39 +0100
Subject: [PATCH 11/19] fix test error

---
 test/mcmc/abstractmcmc.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index e35df1260..3f0a80c9e 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -27,7 +27,7 @@ function initialize_nuts(model::Turing.Model)
     f = Turing.Inference.setvarinfo(
         f,
         DynamicPPL.link!!(Turing.Inference.getvarinfo(f), model),
-        getADType(DynamicPPL.getcontext(f)),
+        getADType(DynamicPPL.getcontext(LogDensityProblemsAD.parent(f))),
     )
 
     # Choose parameter dimensionality and initial parameter value

From ef771a2f088619d3dadf1c824e3f34a339c16446 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 31 Jul 2024 14:40:41 +0100
Subject: [PATCH 12/19] fix more error

---
 src/mcmc/abstractmcmc.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mcmc/abstractmcmc.jl b/src/mcmc/abstractmcmc.jl
index d1fcb4308..a350d2908 100644
--- a/src/mcmc/abstractmcmc.jl
+++ b/src/mcmc/abstractmcmc.jl
@@ -126,7 +126,7 @@ function AbstractMCMC.step(
             varinfo = DynamicPPL.link(varinfo, model)
         end
     end
-    f = setvarinfo(f, varinfo, sampler.alg.adtype)
+    f = setvarinfo(f, varinfo, alg.adtype)
 
     # Then just call `AdvancedHMC.step` with the right arguments.
     if initial_state === nothing

From db78608bb6ab5a755b2835466406643145033076 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 31 Jul 2024 14:46:33 +0100
Subject: [PATCH 13/19] fix error

---
 test/mcmc/abstractmcmc.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index 3f0a80c9e..fd7d9b70c 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -27,7 +27,7 @@ function initialize_nuts(model::Turing.Model)
     f = Turing.Inference.setvarinfo(
         f,
         DynamicPPL.link!!(Turing.Inference.getvarinfo(f), model),
-        getADType(DynamicPPL.getcontext(LogDensityProblemsAD.parent(f))),
+        Turing.getADType(DynamicPPL.getcontext(LogDensityProblemsAD.parent(f))),
     )
 
     # Choose parameter dimensionality and initial parameter value

From 891b00ec7df971df4fe9aee40739c95d9305b9b3 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 31 Jul 2024 15:23:41 +0100
Subject: [PATCH 14/19] fix error

---
 test/mcmc/abstractmcmc.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index fd7d9b70c..5eceb1cb2 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -27,7 +27,7 @@ function initialize_nuts(model::Turing.Model)
     f = Turing.Inference.setvarinfo(
         f,
         DynamicPPL.link!!(Turing.Inference.getvarinfo(f), model),
-        Turing.getADType(DynamicPPL.getcontext(LogDensityProblemsAD.parent(f))),
+        Turing.Inference.getADType(DynamicPPL.getcontext(LogDensityProblemsAD.parent(f))),
     )
 
     # Choose parameter dimensionality and initial parameter value

From fab002dfbca0104c39a331a46ba324509c3ef011 Mon Sep 17 00:00:00 2001
From: Markus Hauru <markus@mhauru.org>
Date: Fri, 30 Aug 2024 14:13:00 +0100
Subject: [PATCH 15/19] Exclude Tapir from AdvancedHMC tests

Co-authored-by: Will Tebbutt <wct23@cam.ac.uk>
---
 test/mcmc/abstractmcmc.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index 5eceb1cb2..a5934d3fd 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -117,7 +117,7 @@ end
 
 @testset "External samplers" begin
     @testset "AdvancedHMC.jl" begin
-        @testset "adtype=$adtype" for adtype in ADUtils.adbackends
+        @testset "adtype=$adtype" for adtype in [AutoForwardDiff(), AutoReverseDiff()]
             @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
                 # Need some functionality to initialize the sampler.
                 # TODO: Remove this once the constructors in the respective packages become "lazy".

From 963269be4ac018926b2290f7e0dff774035a00ab Mon Sep 17 00:00:00 2001
From: Hong Ge <3279477+yebai@users.noreply.github.com>
Date: Mon, 2 Sep 2024 19:09:10 +0100
Subject: [PATCH 16/19] Update ad_utils.jl (#2313)

---
 test/test_utils/ad_utils.jl | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index 90ffb1df4..d6b7fc19b 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -22,14 +22,6 @@ adbackends = [
     Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
 ]
 
-# Tapir isn't supported for older Julia versions, hence the check.
-install_tapir = isdefined(Turing, :AutoTapir)
-if install_tapir
-    # TODO(mhauru) Is there a better way to install optional dependencies like this?
-    Pkg.add("Tapir")
-    push!(adbackends, Turing.AutoTapir(false))
-end
-
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 # Stuff for checking that the right AD backend is being used.
 
@@ -62,6 +54,15 @@ const eltypes_by_adtype = Dict(
     ),
 )
 
+# Tapir isn't supported for older Julia versions, hence the check.
+install_tapir = isdefined(Turing, :AutoTapir)
+if install_tapir
+    # TODO(mhauru) Is there a better way to install optional dependencies like this?
+    Pkg.add("Tapir")
+    push!(adbackends, Turing.AutoTapir(false))
+    push!(eltypes_by_adtype, Turing.AutoTapir => (Tapir.CoDual,))
+end
+
 """
     AbstractWrongADBackendError
 

From 6dac4354bbfcbaaa8818a04a721a523a04105c17 Mon Sep 17 00:00:00 2001
From: Hong Ge <3279477+yebai@users.noreply.github.com>
Date: Mon, 2 Sep 2024 19:14:12 +0100
Subject: [PATCH 17/19] Update test/test_utils/ad_utils.jl

---
 test/test_utils/ad_utils.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index d6b7fc19b..e593ac263 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -59,6 +59,7 @@ install_tapir = isdefined(Turing, :AutoTapir)
 if install_tapir
     # TODO(mhauru) Is there a better way to install optional dependencies like this?
     Pkg.add("Tapir")
+    using Tapir
     push!(adbackends, Turing.AutoTapir(false))
     push!(eltypes_by_adtype, Turing.AutoTapir => (Tapir.CoDual,))
 end

From cb22a05f9075411df61abbbcadf3608d72f2db78 Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Tue, 3 Sep 2024 09:31:54 +0100
Subject: [PATCH 18/19] Move code around in ad_utils.jl

---
 test/test_utils/ad_utils.jl | 40 ++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
index e593ac263..e900a8f69 100644
--- a/test/test_utils/ad_utils.jl
+++ b/test/test_utils/ad_utils.jl
@@ -12,16 +12,6 @@ using Zygote: Zygote
 
 export ADTypeCheckContext, adbackends
 
-# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
-# List of AD backends to test.
-
-"""
-All the ADTypes on which we want to run the tests.
-"""
-adbackends = [
-    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
-]
-
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 # Stuff for checking that the right AD backend is being used.
 
@@ -54,16 +44,6 @@ const eltypes_by_adtype = Dict(
     ),
 )
 
-# Tapir isn't supported for older Julia versions, hence the check.
-install_tapir = isdefined(Turing, :AutoTapir)
-if install_tapir
-    # TODO(mhauru) Is there a better way to install optional dependencies like this?
-    Pkg.add("Tapir")
-    using Tapir
-    push!(adbackends, Turing.AutoTapir(false))
-    push!(eltypes_by_adtype, Turing.AutoTapir => (Tapir.CoDual,))
-end
-
 """
     AbstractWrongADBackendError
 
@@ -294,4 +274,24 @@ Test.@testset "ADTypeCheckContext" begin
     end
 end
 
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+# List of AD backends to test.
+
+"""
+All the ADTypes on which we want to run the tests.
+"""
+adbackends = [
+    Turing.AutoForwardDiff(; chunksize=0), Turing.AutoReverseDiff(; compile=false)
+]
+
+# Tapir isn't supported for older Julia versions, hence the check.
+install_tapir = isdefined(Turing, :AutoTapir)
+if install_tapir
+    # TODO(mhauru) Is there a better way to install optional dependencies like this?
+    Pkg.add("Tapir")
+    using Tapir
+    push!(adbackends, Turing.AutoTapir(false))
+    push!(eltypes_by_adtype, Turing.AutoTapir => (Tapir.CoDual,))
+end
+
 end

From 832552983c2de56e4e4e009c0c60ef7eb251e267 Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Tue, 3 Sep 2024 09:48:54 +0100
Subject: [PATCH 19/19] Add a todo note

---
 test/mcmc/abstractmcmc.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
index a5934d3fd..43e3966a9 100644
--- a/test/mcmc/abstractmcmc.jl
+++ b/test/mcmc/abstractmcmc.jl
@@ -117,6 +117,9 @@ end
 
 @testset "External samplers" begin
     @testset "AdvancedHMC.jl" begin
+        # TODO(mhauru) The below tests fail with Tapir, see
+        # https://github.com/TuringLang/Turing.jl/pull/2289.
+        # Once that is fixed, this should say `for adtype in ADUtils.adbackends`.
         @testset "adtype=$adtype" for adtype in [AutoForwardDiff(), AutoReverseDiff()]
             @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
                 # Need some functionality to initialize the sampler.