diff --git a/.github/workflows/AutoDiffIntegration.yml b/.github/workflows/AutoDiffIntegration.yml
new file mode 100644
index 000000000..3b02354a4
--- /dev/null
+++ b/.github/workflows/AutoDiffIntegration.yml
@@ -0,0 +1,44 @@
+name: Autodiff Integration
+on:
+  workflow_call:
+    inputs:
+      AD:
+        required: true
+        type: string
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  test:
+    name: Julia ${{ matrix.AD }} - ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - 'lts'
+          - '1'
+        os:
+          - ubuntu-latest
+          - macOS-latest
+          - windows-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+      - uses: julia-actions/cache@v1
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          GROUP: AD
+          AD: ${{ inputs.AD }}
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v3
+        with:
+          files: lcov.info
+          
diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml
index de16a19d0..eb97527c2 100644
--- a/.github/workflows/Enzyme.yml
+++ b/.github/workflows/Enzyme.yml
@@ -1,4 +1,4 @@
-name: Enzyme
+name: Enzyme Integration Tests
 on:
   push:
     branches:
@@ -6,35 +6,9 @@ on:
     tags: ['*']
   pull_request:
   workflow_dispatch:
-concurrency:
-  # Skip intermediate builds: always.
-  # Cancel intermediate builds: only if it is a pull request build.
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
-    env: 
-      TEST_GROUP: Enzyme
-    runs-on: ${{ matrix.os }}
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - 'lts'
-          - '1'
-        os:
-          - ubuntu-latest
-          - macOS-latest
-          - windows-latest
-        arch:
-          - x64
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
-      - uses: julia-actions/cache@v1
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
+    uses: ./.github/workflows/AutoDiffIntegration.yml
+    with:
+      AD: Enzyme
diff --git a/.github/workflows/ForwardDiff.yml b/.github/workflows/ForwardDiff.yml
new file mode 100644
index 000000000..830ae9a4e
--- /dev/null
+++ b/.github/workflows/ForwardDiff.yml
@@ -0,0 +1,14 @@
+name: ForwardDiff Integration Tests
+on:
+  push:
+    branches:
+      - main
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  test:
+    uses: ./.github/workflows/AutoDiffIntegration.yml
+    with:
+      AD: ForwardDiff
diff --git a/.github/workflows/Mooncake.yml b/.github/workflows/Mooncake.yml
new file mode 100644
index 000000000..9de098606
--- /dev/null
+++ b/.github/workflows/Mooncake.yml
@@ -0,0 +1,14 @@
+name: Mooncake Integration 
+on:
+  push:
+    branches:
+      - main
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  test:
+    uses: ./.github/workflows/AutoDiffIntegration.yml
+    with:
+      AD: Mooncake
diff --git a/.github/workflows/ReverseDiff.yml b/.github/workflows/ReverseDiff.yml
new file mode 100644
index 000000000..99c9f0d5b
--- /dev/null
+++ b/.github/workflows/ReverseDiff.yml
@@ -0,0 +1,14 @@
+name: ReverseDiff Integration Tests
+on:
+  push:
+    branches:
+      - main
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  test:
+    uses: ./.github/workflows/AutoDiffIntegration.yml
+    with:
+      AD: ReverseDiff
diff --git a/.github/workflows/CI.yml b/.github/workflows/Tests.yml
similarity index 84%
rename from .github/workflows/CI.yml
rename to .github/workflows/Tests.yml
index 06219ec9f..51a6bfeb0 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/Tests.yml
@@ -1,4 +1,4 @@
-name: CI
+name: General Tests
 on:
   push:
     branches:
@@ -13,7 +13,7 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -25,17 +25,17 @@ jobs:
           - ubuntu-latest
           - macOS-latest
           - windows-latest
-        arch:
-          - x64
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        env:
+          GROUP: All
+          AD: ReverseDiff
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v3
         with:
diff --git a/test/algorithms/paramspacesgd/repgradelbo.jl b/test/algorithms/paramspacesgd/repgradelbo.jl
index 2ced719b5..45f22c0ad 100644
--- a/test/algorithms/paramspacesgd/repgradelbo.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo.jl
@@ -1,19 +1,4 @@
 
-AD_repgradelbo_interface = if TEST_GROUP == "Enzyme"
-    [
-        AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    ]
-else
-    [
-        AutoReverseDiff(),
-        AutoZygote(),
-        AutoMooncake(; config=Mooncake.Config()),
-    ]
-end
-
 @testset "interface RepGradELBO" begin
     seed = (0x38bef07cf9cc549d)
     rng = StableRNG(seed)
@@ -25,22 +10,9 @@ end
     q0 = MeanFieldGaussian(zeros(n_dims), Diagonal(ones(n_dims)))
 
     @testset "basic" begin
-        @testset for adtype in AD_repgradelbo_interface, n_montecarlo in [1, 10]
+        @testset for n_montecarlo in [1, 10]
             alg = KLMinRepGradDescent(
-                adtype;
-                n_samples=n_montecarlo,
-                operator=IdentityOperator(),
-                averager=PolynomialAveraging(),
-            )
-            _, info, _ = optimize(rng, alg, 10, model, q0; show_progress=false)
-            @test isfinite(last(info).elbo)
-        end
-    end
-
-    @testset "without mixed ad" begin
-        @testset for adtype in AD_repgradelbo_interface, n_montecarlo in [1, 10]
-            alg = KLMinRepGradDescent(
-                adtype;
+                AD;
                 n_samples=n_montecarlo,
                 operator=IdentityOperator(),
                 averager=PolynomialAveraging(),
@@ -73,7 +45,7 @@ end
     modelstats = normal_meanfield(rng, Float64)
     (; model, μ_true, L_true, n_dims, is_meanfield) = modelstats
 
-    @testset for adtype in AD_repgradelbo_interface, n_montecarlo in [1, 10]
+    @testset for n_montecarlo in [1, 10]
         q_true = MeanFieldGaussian(
             Vector{eltype(μ_true)}(μ_true), Diagonal(Vector{eltype(L_true)}(diag(L_true)))
         )
@@ -81,11 +53,9 @@ end
         obj = RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy())
         out = DiffResults.DiffResult(zero(eltype(params)), similar(params))
 
-        aux = (
-            rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true, adtype=adtype
-        )
+        aux = (rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true, adtype=AD)
         AdvancedVI._value_and_gradient!(
-            AdvancedVI.estimate_repgradelbo_ad_forward, out, adtype, params, aux
+            AdvancedVI.estimate_repgradelbo_ad_forward, out, AD, params, aux
         )
         grad = DiffResults.gradient(out)
         @test norm(grad) ≈ 0 atol = 1e-5
diff --git a/test/algorithms/paramspacesgd/repgradelbo_locationscale.jl b/test/algorithms/paramspacesgd/repgradelbo_locationscale.jl
index 3f7d4f114..9ca41a23e 100644
--- a/test/algorithms/paramspacesgd/repgradelbo_locationscale.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo_locationscale.jl
@@ -1,29 +1,12 @@
-AD_repgradelbo_locationscale = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "inference RepGradELBO VILocationScale" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
-                                                                     [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype)" for realtype in [Float64, Float32],
         (modelname, modelconstr) in
         Dict(:Normal => normal_meanfield, :Normal => normal_fullrank),
         (objname, objective) in Dict(
             :RepGradELBOClosedFormEntropy => RepGradELBO(10),
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(10; entropy=StickingTheLandingEntropy()),
-        ),
-        (adbackname, adtype) in AD_repgradelbo_locationscale
+        )
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -33,7 +16,7 @@ end
 
         T = 1000
         η = 1e-3
-        alg = KLMinRepGradDescent(adtype; optimizer=Descent(η))
+        alg = KLMinRepGradDescent(AD; optimizer=Descent(η))
 
         q0 = if is_meanfield
             MeanFieldGaussian(zeros(realtype, n_dims), Diagonal(ones(realtype, n_dims)))
diff --git a/test/algorithms/paramspacesgd/repgradelbo_locationscale_bijectors.jl b/test/algorithms/paramspacesgd/repgradelbo_locationscale_bijectors.jl
index 33995ee84..b96f520c4 100644
--- a/test/algorithms/paramspacesgd/repgradelbo_locationscale_bijectors.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo_locationscale_bijectors.jl
@@ -1,29 +1,12 @@
-AD_repgradelbo_locationscale_bijectors = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "inference RepGradELBO VILocationScale Bijectors" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
-                                                                     [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype)" for realtype in [Float64, Float32],
         (modelname, modelconstr) in
         Dict(:NormalLogNormalMeanField => normallognormal_meanfield),
         (objname, objective) in Dict(
             :RepGradELBOClosedFormEntropy => RepGradELBO(10),
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(10; entropy=StickingTheLandingEntropy()),
-        ),
-        (adbackname, adtype) in AD_repgradelbo_locationscale_bijectors
+        )
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -33,7 +16,7 @@ end
 
         T = 1000
         η = 1e-3
-        alg = KLMinRepGradDescent(adtype; optimizer=Descent(η))
+        alg = KLMinRepGradDescent(AD; optimizer=Descent(η))
 
         b = Bijectors.bijector(model)
         b⁻¹ = inverse(b)
diff --git a/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale.jl b/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale.jl
index 624a292f7..1eddac5f1 100644
--- a/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale.jl
@@ -1,22 +1,5 @@
-
-AD_repgradelbo_locationscale = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "inference RepGradELBO Proximal VILocationScale" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
-                                                                     [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype)" for realtype in [Float64, Float32],
         (modelname, modelconstr) in
         Dict(:Normal => normal_meanfield, :Normal => normal_fullrank),
         (objname, objective) in Dict(
@@ -24,8 +7,7 @@ end
                 RepGradELBO(10; entropy=ClosedFormEntropyZeroGradient()),
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(10; entropy=StickingTheLandingEntropyZeroGradient()),
-        ),
-        (adbackname, adtype) in AD_repgradelbo_locationscale
+        )
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -42,7 +24,7 @@ end
 
         T = 1000
         η = 1e-3
-        alg = KLMinRepGradProxDescent(adtype; optimizer=Descent(η))
+        alg = KLMinRepGradProxDescent(AD; optimizer=Descent(η))
 
         # For small enough η, the error of SGD, Δλ, is bounded as
         #     Δλ ≤ ρ^T Δλ0 + O(η),
diff --git a/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale_bijectors.jl b/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale_bijectors.jl
index dcd9722a8..2875c1c72 100644
--- a/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale_bijectors.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale_bijectors.jl
@@ -1,21 +1,5 @@
-AD_repgradelbo_locationscale_bijectors = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "inference RepGradELBO Proximal VILocationScale Bijectors" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
-                                                                     [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype)" for realtype in [Float64, Float32],
         (modelname, modelconstr) in
         Dict(:NormalLogNormalMeanField => normallognormal_meanfield),
         (objname, objective) in Dict(
@@ -23,8 +7,7 @@ end
                 RepGradELBO(10; entropy=ClosedFormEntropyZeroGradient()),
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(10; entropy=StickingTheLandingEntropyZeroGradient()),
-        ),
-        (adbackname, adtype) in AD_repgradelbo_locationscale_bijectors
+        )
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -34,7 +17,7 @@ end
 
         T = 1000
         η = 1e-3
-        alg = KLMinRepGradProxDescent(adtype; optimizer=Descent(η))
+        alg = KLMinRepGradProxDescent(AD; optimizer=Descent(η))
 
         b = Bijectors.bijector(model)
         b⁻¹ = inverse(b)
diff --git a/test/algorithms/paramspacesgd/scoregradelbo.jl b/test/algorithms/paramspacesgd/scoregradelbo.jl
index 011aea6a1..cd21de53c 100644
--- a/test/algorithms/paramspacesgd/scoregradelbo.jl
+++ b/test/algorithms/paramspacesgd/scoregradelbo.jl
@@ -1,15 +1,3 @@
-
-AD_scoregradelbo_interface = if TEST_GROUP == "Enzyme"
-    [AutoEnzyme()]
-else
-    [
-        AutoForwardDiff(),
-        AutoReverseDiff(),
-        AutoZygote(),
-        AutoMooncake(; config=Mooncake.Config()),
-    ]
-end
-
 @testset "interface ScoreGradELBO" begin
     seed = (0x38bef07cf9cc549d)
     rng = StableRNG(seed)
@@ -21,10 +9,8 @@ end
     q0 = MeanFieldGaussian(zeros(n_dims), Diagonal(ones(n_dims)))
 
     @testset "basic" begin
-        @testset for adtype in AD_scoregradelbo_interface, n_montecarlo in [1, 10]
-            alg = KLMinScoreGradDescent(
-                adtype; n_samples=n_montecarlo, optimizer=Descent(1e-5)
-            )
+        @testset for n_montecarlo in [1, 10]
+            alg = KLMinScoreGradDescent(AD; n_samples=n_montecarlo, optimizer=Descent(1e-5))
             _, info, _ = optimize(rng, alg, 10, model, q0; show_progress=false)
             @assert isfinite(last(info).elbo)
         end
diff --git a/test/algorithms/paramspacesgd/scoregradelbo_locationscale.jl b/test/algorithms/paramspacesgd/scoregradelbo_locationscale.jl
index 2505a3669..109a2b044 100644
--- a/test/algorithms/paramspacesgd/scoregradelbo_locationscale.jl
+++ b/test/algorithms/paramspacesgd/scoregradelbo_locationscale.jl
@@ -1,25 +1,7 @@
-
-AD_scoregradelbo_locationscale = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ForwarDiff => AutoForwardDiff(),
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "inference ScoreGradELBO VILocationScale" begin
-    @testset "$(modelname) $(realtype) $(adbackname)" for realtype in [Float64, Float32],
+    @testset "$(modelname) $(realtype)" for realtype in [Float64, Float32],
         (modelname, modelconstr) in
-        Dict(:Normal => normal_meanfield, :Normal => normal_fullrank),
-        (adbackname, adtype) in AD_scoregradelbo_locationscale
+        Dict(:Normal => normal_meanfield, :Normal => normal_fullrank)
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -30,7 +12,7 @@ end
         T = 1000
         η = 1e-4
         opt = Optimisers.Descent(η)
-        alg = KLMinScoreGradDescent(adtype; n_samples=10, optimizer=opt)
+        alg = KLMinScoreGradDescent(AD; n_samples=10, optimizer=opt)
 
         q0 = if is_meanfield
             MeanFieldGaussian(zeros(realtype, n_dims), Diagonal(ones(realtype, n_dims)))
diff --git a/test/algorithms/paramspacesgd/scoregradelbo_locationscale_bijectors.jl b/test/algorithms/paramspacesgd/scoregradelbo_locationscale_bijectors.jl
index 63add9269..84d570f2a 100644
--- a/test/algorithms/paramspacesgd/scoregradelbo_locationscale_bijectors.jl
+++ b/test/algorithms/paramspacesgd/scoregradelbo_locationscale_bijectors.jl
@@ -1,24 +1,7 @@
-AD_scoregradelbo_locationscale_bijectors = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ForwarDiff => AutoForwardDiff(),
-        :ReverseDiff => AutoReverseDiff(),
-        #:Zygote => AutoZygote(),
-        #:Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "inference ScoreGradELBO VILocationScale Bijectors" begin
-    @testset "$(modelname) $(realtype) $(adbackname)" for realtype in [Float64, Float32],
+    @testset "$(modelname) $(realtype)" for realtype in [Float64, Float32],
         (modelname, modelconstr) in
-        Dict(:NormalLogNormalMeanField => normallognormal_meanfield),
-        (adbackname, adtype) in AD_scoregradelbo_locationscale_bijectors
+        Dict(:NormalLogNormalMeanField => normallognormal_meanfield)
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -29,7 +12,7 @@ end
         T = 1000
         η = 1e-4
         opt = Optimisers.Descent(η)
-        alg = KLMinScoreGradDescent(adtype; n_samples=10, optimizer=opt)
+        alg = KLMinScoreGradDescent(AD; n_samples=10, optimizer=opt)
 
         b = Bijectors.bijector(model)
         b⁻¹ = inverse(b)
diff --git a/test/general/ad.jl b/test/general/ad.jl
index ac9a82bfd..251eead5e 100644
--- a/test/general/ad.jl
+++ b/test/general/ad.jl
@@ -1,36 +1,20 @@
 
-AD_interface = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ForwarDiff => AutoForwardDiff(),
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 @testset "ad" begin
-    @testset "$(adname)" for (adname, adtype) in AD_interface
+    @testset "value_and_gradient!" begin
         D = 10
         A = randn(D, D)
         λ = randn(D)
         b = randn(D)
         grad_buf = DiffResults.GradientResult(λ)
         f(λ′, aux) = λ′' * A * λ′ / 2 + dot(aux.b, λ′)
-        AdvancedVI._value_and_gradient!(f, grad_buf, adtype, λ, (b=b,))
+        AdvancedVI._value_and_gradient!(f, grad_buf, AD, λ, (b=b,))
         ∇ = DiffResults.gradient(grad_buf)
         f = DiffResults.value(grad_buf)
         @test ∇ ≈ (A + A') * λ / 2 + b
         @test f ≈ λ' * A * λ / 2 + dot(b, λ)
     end
 
-    @testset "$(adname) with prep" for (adname, adtype) in AD_interface
+    @testset "value_and_gradient! with prep" begin
         D = 10
         λ = randn(D)
         A = randn(D, D)
@@ -38,10 +22,10 @@ end
 
         b_prep = randn(D)
         f(λ′, aux) = λ′' * A * λ′ / 2 + dot(aux.b, λ′)
-        prep = AdvancedVI._prepare_gradient(f, adtype, λ, (b=b_prep,))
+        prep = AdvancedVI._prepare_gradient(f, AD, λ, (b=b_prep,))
 
         b = randn(D)
-        AdvancedVI._value_and_gradient!(f, grad_buf, prep, adtype, λ, (b=b,))
+        AdvancedVI._value_and_gradient!(f, grad_buf, prep, AD, λ, (b=b,))
 
         ∇ = DiffResults.gradient(grad_buf)
         f = DiffResults.value(grad_buf)
diff --git a/test/general/mixedad_logdensity.jl b/test/general/mixedad_logdensity.jl
index 2f9edaf52..bee48a397 100644
--- a/test/general/mixedad_logdensity.jl
+++ b/test/general/mixedad_logdensity.jl
@@ -1,19 +1,4 @@
 
-AD_mixedad = if TEST_GROUP == "Enzyme"
-    Dict(
-        :Enzyme => AutoEnzyme(;
-            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-            function_annotation=Enzyme.Const,
-        ),
-    )
-else
-    Dict(
-        :ReverseDiff => AutoReverseDiff(),
-        :Zygote => AutoZygote(),
-        :Mooncake => AutoMooncake(; config=Mooncake.Config()),
-    )
-end
-
 struct MixedADTestModel end
 
 function LogDensityProblems.logdensity(::MixedADTestModel, θ)
@@ -46,27 +31,32 @@ function mixedad_test_fwd(x, prob)
     )/2
 end
 
-@testset "MixedADLogDensityProblem" begin
-    model = MixedADTestModel()
-    model_ad = AdvancedVI.MixedADLogDensityProblem(model)
-
-    d = 3
-    x = ones(Float64, d)
-
-    @testset "interface" begin
-        @test LogDensityProblems.dimension(model) == LogDensityProblems.dimension(model_ad)
-        @test last(LogDensityProblems.logdensity(model, x)) ≈
-            last(LogDensityProblems.logdensity(model_ad, x))
-    end
-
-    @testset "rrule under $(adname)" for (adname, adtype) in AD_mixedad
-        out = DiffResults.DiffResult(0.0, zeros(d))
-        AdvancedVI._value_and_gradient!(mixedad_test_fwd, out, adtype, x, model_ad)
-        @test DiffResults.gradient(out) ≈ EXPECTED_RESULT
-
-        out = DiffResults.DiffResult(0.0, zeros(d))
-        prep = AdvancedVI._prepare_gradient(mixedad_test_fwd, adtype, x, model_ad)
-        AdvancedVI._value_and_gradient!(mixedad_test_fwd, out, prep, adtype, x, model_ad)
-        @test DiffResults.gradient(out) ≈ EXPECTED_RESULT
+# MixedADLogDensityProblem only supports ReverseDiff, Zygote, Enzyme, Mooncake in reverse-mode
+if (AD isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoEnzyme,<:AutoMooncake}) &&
+    (ADTypes.mode(AD) isa ADTypes.ReverseMode)
+    @testset "MixedADLogDensityProblem" begin
+        model = MixedADTestModel()
+        model_ad = AdvancedVI.MixedADLogDensityProblem(model)
+
+        d = 3
+        x = ones(Float64, d)
+
+        @testset "interface" begin
+            @test LogDensityProblems.dimension(model) ==
+                LogDensityProblems.dimension(model_ad)
+            @test last(LogDensityProblems.logdensity(model, x)) ≈
+                last(LogDensityProblems.logdensity(model_ad, x))
+        end
+
+        @testset "custom rrule" begin
+            out = DiffResults.DiffResult(0.0, zeros(d))
+            AdvancedVI._value_and_gradient!(mixedad_test_fwd, out, AD, x, model_ad)
+            @test DiffResults.gradient(out) ≈ EXPECTED_RESULT
+
+            out = DiffResults.DiffResult(0.0, zeros(d))
+            prep = AdvancedVI._prepare_gradient(mixedad_test_fwd, AD, x, model_ad)
+            AdvancedVI._value_and_gradient!(mixedad_test_fwd, out, prep, AD, x, model_ad)
+            @test DiffResults.gradient(out) ≈ EXPECTED_RESULT
+        end
     end
 end
diff --git a/test/general/optimize.jl b/test/general/optimize.jl
index 5849e2bc2..71c3e4fb4 100644
--- a/test/general/optimize.jl
+++ b/test/general/optimize.jl
@@ -11,11 +11,10 @@
     q0 = MeanFieldGaussian(zeros(Float64, n_dims), Diagonal(ones(Float64, n_dims)))
     obj = RepGradELBO(10)
 
-    adtype = AutoReverseDiff()
     optimizer = Optimisers.Adam(1e-2)
     averager = PolynomialAveraging()
 
-    alg = ParamSpaceSGD(obj, adtype, optimizer, averager, IdentityOperator())
+    alg = ParamSpaceSGD(obj, AD, optimizer, averager, IdentityOperator())
 
     @testset "default_rng" begin
         optimize(alg, T, model, q0; show_progress=false)
diff --git a/test/general/proximal_location_scale_entropy.jl b/test/general/proximal_location_scale_entropy.jl
index ddbbf7300..fb58535c6 100644
--- a/test/general/proximal_location_scale_entropy.jl
+++ b/test/general/proximal_location_scale_entropy.jl
@@ -51,9 +51,11 @@
             q′ = re(params′)
             scale′ = isnothing(bijector) ? q′.scale : q′.dist.scale
 
-            grad_left = only(Zygote.gradient(L_ -> first(logabsdet(L_)), scale′))
-            grad_right = only(
-                Zygote.gradient(L_ -> sum(abs2, L_ - L) / (2 * stepsize), scale′)
+            grad_left = ReverseDiff.gradient(
+                L_ -> first(logabsdet(LowerTriangular(reshape(L_, d, d)))), vec(scale′)
+            )
+            grad_right = ReverseDiff.gradient(
+                L_ -> sum(abs2, reshape(L_, d, d) - L) / (2 * stepsize), vec(scale′)
             )
 
             @test grad_left ≈ grad_right
diff --git a/test/runtests.jl b/test/runtests.jl
index e72e00af8..3e33b7405 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,11 +2,13 @@
 using Test
 using Test: @testset, @test
 
+using ADTypes
 using Base.Iterators
 using Bijectors
 using DiffResults
 using Distributions
 using FillArrays
+using ForwardDiff, ReverseDiff
 using LinearAlgebra
 using LogDensityProblems
 using Optimisers
@@ -16,16 +18,28 @@ using Random, StableRNGs
 using Statistics
 using StatsBase
 
-using ADTypes
-using ForwardDiff, ReverseDiff, Zygote, Mooncake
-
 using AdvancedVI
 
 const PROGRESS = haskey(ENV, "PROGRESS")
-const TEST_GROUP = get(ENV, "TEST_GROUP", "All")
+const GROUP = get(ENV, "GROUP", "All")
+const AD_str = get(ENV, "AD", "ReverseDiff")
 
-if TEST_GROUP == "Enzyme"
+const AD = if AD_str == "ReverseDiff"
+    AutoReverseDiff()
+elseif AD_str == "ForwardDiff"
+    AutoForwardDiff()
+elseif AD_str == "Zygote"
+    using Zygote
+    AutoZygote()
+elseif AD_str == "Mooncake"
+    using Mooncake
+    AutoMooncake(; config=Mooncake.Config())
+elseif AD_str == "Enzyme"
     using Enzyme
+    AutoEnzyme(;
+        mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+        function_annotation=Enzyme.Const,
+    )
 end
 
 # Models for Inference Tests
@@ -40,27 +54,23 @@ end
 include("models/normal.jl")
 include("models/normallognormal.jl")
 
-if TEST_GROUP == "All" || TEST_GROUP == "General"
-    # Interface tests that do not involve testing on Enzyme
+if GROUP == "All" || GROUP == "GENERAL"
+    # Tests that do not need to check correct integration with AD backends
     include("general/optimize.jl")
+    include("general/proximal_location_scale_entropy.jl")
     include("general/rules.jl")
     include("general/averaging.jl")
     include("general/clip_scale.jl")
-    include("general/proximal_location_scale_entropy.jl")
-    include("general/mixedad_logdensity.jl")
-end
 
-if TEST_GROUP == "All" || TEST_GROUP == "General" || TEST_GROUP == "Enzyme"
-    # Interface tests that involve testing on Enzyme
-    include("general/ad.jl")
-end
-
-if TEST_GROUP == "All" || TEST_GROUP == "Families"
     include("families/location_scale.jl")
     include("families/location_scale_low_rank.jl")
 end
 
-if TEST_GROUP == "All" || TEST_GROUP == "ParamSpaceSGD" || TEST_GROUP == "Enzyme"
+if GROUP == "All" || GROUP == "AD"
+    # Tests that need to check correctness of the integration with AD backends
+    include("general/ad.jl")
+    include("general/mixedad_logdensity.jl")
+
     include("algorithms/paramspacesgd/repgradelbo.jl")
     include("algorithms/paramspacesgd/scoregradelbo.jl")
     include("algorithms/paramspacesgd/repgradelbo_locationscale.jl")