From 347e7f9138077d0ea108b6a0ff174a58a0baf40a Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 16:26:06 -0400
Subject: [PATCH 01/21] refactor split AD test loops into workflows

---
 .github/workflows/AutoDiffIntegration.yml     | 48 +++++++++++++++++++
 .github/workflows/Enzyme.yml                  | 38 +++------------
 .github/workflows/ForwardDiff.yml             | 14 ++++++
 .github/workflows/Mooncake.yml                | 14 ++++++
 .github/workflows/ReverseDiff.yml             | 14 ++++++
 .github/workflows/{CI.yml => Tests.yml}       |  5 +-
 .../repgradelbo_locationscale.jl              | 23 ++-------
 .../repgradelbo_locationscale_bijectors.jl    | 23 ++-------
 .../repgradelbo_proximal_locationscale.jl     | 24 ++--------
 ...adelbo_proximal_locationscale_bijectors.jl | 23 ++-------
 .../algorithms/paramspacesgd/scoregradelbo.jl | 18 +------
 .../scoregradelbo_locationscale.jl            | 24 ++--------
 .../scoregradelbo_locationscale_bijectors.jl  | 23 ++-------
 test/general/ad.jl                            | 26 ++--------
 test/general/optimize.jl                      |  3 +-
 15 files changed, 126 insertions(+), 194 deletions(-)
 create mode 100644 .github/workflows/AutoDiffIntegration.yml
 create mode 100644 .github/workflows/ForwardDiff.yml
 create mode 100644 .github/workflows/Mooncake.yml
 create mode 100644 .github/workflows/ReverseDiff.yml
 rename .github/workflows/{CI.yml => Tests.yml} (94%)

diff --git a/.github/workflows/AutoDiffIntegration.yml b/.github/workflows/AutoDiffIntegration.yml
new file mode 100644
index 000000000..39c3924af
--- /dev/null
+++ b/.github/workflows/AutoDiffIntegration.yml
@@ -0,0 +1,48 @@
+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 }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    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
+        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..a0c5b0a84 100644
--- a/.github/workflows/Enzyme.yml
+++ b/.github/workflows/Enzyme.yml
@@ -1,4 +1,4 @@
-name: Enzyme
+name: Enzyme Integration 
 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
+  Enzyme:
+    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..da2e6f02c
--- /dev/null
+++ b/.github/workflows/ForwardDiff.yml
@@ -0,0 +1,14 @@
+name: ForwardDiff Integration 
+on:
+  push:
+    branches:
+      - main
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  Enzyme:
+    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..d7b2e8f8c
--- /dev/null
+++ b/.github/workflows/Mooncake.yml
@@ -0,0 +1,14 @@
+name: Mooncake Integration 
+on:
+  push:
+    branches:
+      - main
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  Enzyme:
+    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..705c0eaca
--- /dev/null
+++ b/.github/workflows/ReverseDiff.yml
@@ -0,0 +1,14 @@
+name: ReverseDiff Integration 
+on:
+  push:
+    branches:
+      - main
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  Enzyme:
+    uses: ./.github/workflows/AutoDiffIntegration.yml
+    with:
+      AD: ReverseDiff
diff --git a/.github/workflows/CI.yml b/.github/workflows/Tests.yml
similarity index 94%
rename from .github/workflows/CI.yml
rename to .github/workflows/Tests.yml
index 06219ec9f..10a279dd1 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/Tests.yml
@@ -1,4 +1,4 @@
-name: CI
+name: Tests
 on:
   push:
     branches:
@@ -36,6 +36,9 @@ jobs:
       - 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_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/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)

From 09993415620daacc27687a301d4d86426960ac84 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:05:09 -0400
Subject: [PATCH 02/21] fix repgradelbo interface test to match new workflow
 interface

---
 test/algorithms/paramspacesgd/repgradelbo.jl | 23 ++++----------------
 1 file changed, 4 insertions(+), 19 deletions(-)

diff --git a/test/algorithms/paramspacesgd/repgradelbo.jl b/test/algorithms/paramspacesgd/repgradelbo.jl
index 2ced719b5..e429b1077 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,9 +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;
+                AD;
                 n_samples=n_montecarlo,
                 operator=IdentityOperator(),
                 averager=PolynomialAveraging(),
@@ -38,9 +23,9 @@ end
     end
 
     @testset "without mixed ad" begin
-        @testset for adtype in AD_repgradelbo_interface, n_montecarlo in [1, 10]
+        @testset for n_montecarlo in [1, 10]
             alg = KLMinRepGradDescent(
-                adtype;
+                AD;
                 n_samples=n_montecarlo,
                 operator=IdentityOperator(),
                 averager=PolynomialAveraging(),

From 1f6383922dc47ff68de2af93322d51d4dd3657b4 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:21:11 -0400
Subject: [PATCH 03/21] fix missing detail in info message for mixed ad
 repgradelbo

---
 src/algorithms/paramspacesgd/repgradelbo.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/algorithms/paramspacesgd/repgradelbo.jl b/src/algorithms/paramspacesgd/repgradelbo.jl
index 371f119f7..210c84672 100644
--- a/src/algorithms/paramspacesgd/repgradelbo.jl
+++ b/src/algorithms/paramspacesgd/repgradelbo.jl
@@ -41,7 +41,7 @@ function init(
             adtype isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoMooncake,<:AutoEnzyme} &&
             ADTypes.mode(adtype) == ADTypes.ReverseMode
         )
-            @info "The supplied target `LogDensityProblem` has a differentiation capability of $(capability) >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoZygote`, `AutoMooncake`, or `AutoEnzyme{Enzyme.Reverse,<:Any}`."
+            @info "The supplied target `LogDensityProblem` has a differentiation capability of $(capability) >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoZygote`, `AutoMooncake`, or `AutoEnzyme` in reverse mode."
         end
         MixedADLogDensityProblem(prob)
     end

From 188117ea8c4367e2f6789ce1252ee2c95d3869a6 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:25:22 -0400
Subject: [PATCH 04/21] refactor make info string for mixed ad repgradelbo more
 consistent

---
 src/algorithms/paramspacesgd/repgradelbo.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/algorithms/paramspacesgd/repgradelbo.jl b/src/algorithms/paramspacesgd/repgradelbo.jl
index 210c84672..546a0139b 100644
--- a/src/algorithms/paramspacesgd/repgradelbo.jl
+++ b/src/algorithms/paramspacesgd/repgradelbo.jl
@@ -34,14 +34,14 @@ function init(
     q_stop = restructure(params)
     capability = LogDensityProblems.capabilities(Prob)
     ad_prob = if capability < LogDensityProblems.LogDensityOrder{1}()
-        @info "The capability of the provided log-density problem $(capability) is less than $(LogDensityProblems.LogDensityOrder{1}()). `AdvancedVI` will attempt to directly differentiate through `LogDensityProblems.logdensity`. If this is not intended, please supply a log-density problem with capability at least $(LogDensityProblems.LogDensityOrder{1}())"
+        @info "The capability of the supplied `LogDensityProblem` $(capability) is less than $(LogDensityProblems.LogDensityOrder{1}()). `AdvancedVI` will attempt to directly differentiate through `LogDensityProblems.logdensity`. If this is not intended, please supply a log-density problem with capability at least $(LogDensityProblems.LogDensityOrder{1}())"
         prob
     else
         if !(
             adtype isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoMooncake,<:AutoEnzyme} &&
             ADTypes.mode(adtype) == ADTypes.ReverseMode
         )
-            @info "The supplied target `LogDensityProblem` has a differentiation capability of $(capability) >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoZygote`, `AutoMooncake`, or `AutoEnzyme` in reverse mode."
+            @info "The capability of the supplied target `LogDensityProblem` $(capability) is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoZygote`, `AutoMooncake`, or `AutoEnzyme` in reverse mode."
         end
         MixedADLogDensityProblem(prob)
     end

From 02bd74b4a6e1148d633f8e87007960af5daa6775 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:28:48 -0400
Subject: [PATCH 05/21] fix bug in mixed ad repgradelbo adtype condition check

---
 src/algorithms/paramspacesgd/repgradelbo.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/algorithms/paramspacesgd/repgradelbo.jl b/src/algorithms/paramspacesgd/repgradelbo.jl
index 546a0139b..cb9dd3982 100644
--- a/src/algorithms/paramspacesgd/repgradelbo.jl
+++ b/src/algorithms/paramspacesgd/repgradelbo.jl
@@ -39,7 +39,7 @@ function init(
     else
         if !(
             adtype isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoMooncake,<:AutoEnzyme} &&
-            ADTypes.mode(adtype) == ADTypes.ReverseMode
+            ADTypes.mode(adtype) isa ADTypes.ReverseMode
         )
             @info "The capability of the supplied target `LogDensityProblem` $(capability) is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoZygote`, `AutoMooncake`, or `AutoEnzyme` in reverse mode."
         end

From bc67eb9a3335940c2251ab0f4f130a0335fc80ce Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:30:49 -0400
Subject: [PATCH 06/21] fix update adtype selection to match new workflow-based
 system

---
 test/algorithms/paramspacesgd/repgradelbo.jl  |  6 ++--
 test/general/mixedad_logdensity.jl            | 23 +++---------
 .../proximal_location_scale_entropy.jl        |  8 +++--
 test/runtests.jl                              | 35 +++++++++++--------
 4 files changed, 33 insertions(+), 39 deletions(-)

diff --git a/test/algorithms/paramspacesgd/repgradelbo.jl b/test/algorithms/paramspacesgd/repgradelbo.jl
index e429b1077..bfe92cf16 100644
--- a/test/algorithms/paramspacesgd/repgradelbo.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo.jl
@@ -58,7 +58,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)))
         )
@@ -67,10 +67,10 @@ end
         out = DiffResults.DiffResult(zero(eltype(params)), similar(params))
 
         aux = (
-            rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true, adtype=adtype
+            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/general/mixedad_logdensity.jl b/test/general/mixedad_logdensity.jl
index 2f9edaf52..cc7d0f89f 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, θ)
@@ -59,14 +44,14 @@ end
             last(LogDensityProblems.logdensity(model_ad, x))
     end
 
-    @testset "rrule under $(adname)" for (adname, adtype) in AD_mixedad
+    @testset "custom rrule" begin
         out = DiffResults.DiffResult(0.0, zeros(d))
-        AdvancedVI._value_and_gradient!(mixedad_test_fwd, out, adtype, x, model_ad)
+        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, adtype, x, model_ad)
-        AdvancedVI._value_and_gradient!(mixedad_test_fwd, out, prep, adtype, x, model_ad)
+        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
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..07e70a1f0 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,26 @@ 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 == "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,8 +52,8 @@ 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"
+    include("general/ad.jl")
     include("general/optimize.jl")
     include("general/rules.jl")
     include("general/averaging.jl")
@@ -50,17 +62,12 @@ if TEST_GROUP == "All" || TEST_GROUP == "General"
     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"
+if GROUP == "All" || GROUP == "General" || 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 == "ParamSpaceSGD"
     include("algorithms/paramspacesgd/repgradelbo.jl")
     include("algorithms/paramspacesgd/scoregradelbo.jl")
     include("algorithms/paramspacesgd/repgradelbo_locationscale.jl")

From 8479b085cf0722dbade7a5f52cc95c69bb21d08a Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:35:54 -0400
Subject: [PATCH 07/21] fix make reverse diff the default backend and remove
 its workflow

---
 .github/workflows/ReverseDiff.yml | 14 --------------
 .github/workflows/Tests.yml       |  2 +-
 2 files changed, 1 insertion(+), 15 deletions(-)
 delete mode 100644 .github/workflows/ReverseDiff.yml

diff --git a/.github/workflows/ReverseDiff.yml b/.github/workflows/ReverseDiff.yml
deleted file mode 100644
index 705c0eaca..000000000
--- a/.github/workflows/ReverseDiff.yml
+++ /dev/null
@@ -1,14 +0,0 @@
-name: ReverseDiff Integration 
-on:
-  push:
-    branches:
-      - main
-    tags: ['*']
-  pull_request:
-  workflow_dispatch:
-
-jobs:
-  Enzyme:
-    uses: ./.github/workflows/AutoDiffIntegration.yml
-    with:
-      AD: ReverseDiff
diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml
index 10a279dd1..2a992a4c9 100644
--- a/.github/workflows/Tests.yml
+++ b/.github/workflows/Tests.yml
@@ -1,4 +1,4 @@
-name: Tests
+name: All Tests (ReverseDiff)
 on:
   push:
     branches:

From db6afa965cfebc4c1f847e48019867677555afbc Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:36:26 -0400
Subject: [PATCH 08/21] fix run formatter

---
 test/algorithms/paramspacesgd/repgradelbo.jl | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/test/algorithms/paramspacesgd/repgradelbo.jl b/test/algorithms/paramspacesgd/repgradelbo.jl
index bfe92cf16..48742db50 100644
--- a/test/algorithms/paramspacesgd/repgradelbo.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo.jl
@@ -66,9 +66,7 @@ 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=AD
-        )
+        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, AD, params, aux
         )

From 13a6e96a7ce2521b3b175173389ca50ffdfd8919 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 17:44:16 -0400
Subject: [PATCH 09/21] fix improve workflow name for All Tests workflow

---
 .github/workflows/Tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml
index 2a992a4c9..837f35200 100644
--- a/.github/workflows/Tests.yml
+++ b/.github/workflows/Tests.yml
@@ -1,4 +1,4 @@
-name: All Tests (ReverseDiff)
+name: All Tests with ReverseDiff Integration
 on:
   push:
     branches:

From 864f5e94552424aec2719e3c55246a4cc701df3b Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 18:25:18 -0400
Subject: [PATCH 10/21] fix turn on tests for "AD" test group

---
 test/runtests.jl | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 07e70a1f0..9ebf1032e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -53,13 +53,15 @@ include("models/normal.jl")
 include("models/normallognormal.jl")
 
 if GROUP == "All" || GROUP == "General"
-    include("general/ad.jl")
-    include("general/optimize.jl")
+    if  GROUP == "AD"
+        include("general/ad.jl")
+        include("general/optimize.jl")
+        include("general/proximal_location_scale_entropy.jl")
+        include("general/mixedad_logdensity.jl")
+    end
     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 GROUP == "All" || GROUP == "General" || GROUP == "Families"
@@ -67,7 +69,7 @@ if GROUP == "All" || GROUP == "General" || GROUP == "Families"
     include("families/location_scale_low_rank.jl")
 end
 
-if GROUP == "All" || GROUP == "ParamSpaceSGD"
+if GROUP == "All" || GROUP == "ParamSpaceSGD" || GROUP == "AD"
     include("algorithms/paramspacesgd/repgradelbo.jl")
     include("algorithms/paramspacesgd/scoregradelbo.jl")
     include("algorithms/paramspacesgd/repgradelbo_locationscale.jl")

From fd12f62e9262f8144589695f430866930d1c13e6 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 19:00:36 -0400
Subject: [PATCH 11/21] fix name of jobs of AD integration tests

---
 .github/workflows/Enzyme.yml      | 2 +-
 .github/workflows/ForwardDiff.yml | 2 +-
 .github/workflows/Mooncake.yml    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml
index a0c5b0a84..cf6555f30 100644
--- a/.github/workflows/Enzyme.yml
+++ b/.github/workflows/Enzyme.yml
@@ -8,7 +8,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  Enzyme:
+  test:
     uses: ./.github/workflows/AutoDiffIntegration.yml
     with:
       AD: Enzyme
diff --git a/.github/workflows/ForwardDiff.yml b/.github/workflows/ForwardDiff.yml
index da2e6f02c..2100d7472 100644
--- a/.github/workflows/ForwardDiff.yml
+++ b/.github/workflows/ForwardDiff.yml
@@ -8,7 +8,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  Enzyme:
+  Tests:
     uses: ./.github/workflows/AutoDiffIntegration.yml
     with:
       AD: ForwardDiff
diff --git a/.github/workflows/Mooncake.yml b/.github/workflows/Mooncake.yml
index d7b2e8f8c..9de098606 100644
--- a/.github/workflows/Mooncake.yml
+++ b/.github/workflows/Mooncake.yml
@@ -8,7 +8,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  Enzyme:
+  test:
     uses: ./.github/workflows/AutoDiffIntegration.yml
     with:
       AD: Mooncake

From 6ea485fd4666417a529be5165ff89290f587a760 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 19 Aug 2025 19:01:28 -0400
Subject: [PATCH 12/21] add missing forwarddiff AD string translation in tests

---
 test/runtests.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/test/runtests.jl b/test/runtests.jl
index 9ebf1032e..df79081d1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -26,6 +26,8 @@ const AD_str = get(ENV, "AD", "ReverseDiff")
 
 const AD = if AD_str == "ReverseDiff"
     AutoReverseDiff()
+elseif AD_str == "ForwardDiff"
+    AutoForwardDiff()
 elseif AD_str == "Zygote"
     using Zygote
     AutoZygote()

From e843e57628bec6d41d8bc67d141fd5d7177df791 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <msca8h@naver.com>
Date: Wed, 20 Aug 2025 14:04:28 -0400
Subject: [PATCH 13/21] run formatter

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index df79081d1..9e5716c0f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -55,7 +55,7 @@ include("models/normal.jl")
 include("models/normallognormal.jl")
 
 if GROUP == "All" || GROUP == "General"
-    if  GROUP == "AD"
+    if GROUP == "AD"
         include("general/ad.jl")
         include("general/optimize.jl")
         include("general/proximal_location_scale_entropy.jl")

From b5484206a14054bb9006097de0f6606ae0e29589 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Thu, 21 Aug 2025 12:02:15 -0400
Subject: [PATCH 14/21] refactor test workflows don't specify architecture

---
 .github/workflows/AutoDiffIntegration.yml | 5 +----
 .github/workflows/ForwardDiff.yml         | 2 +-
 .github/workflows/Tests.yml               | 5 +----
 3 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/AutoDiffIntegration.yml b/.github/workflows/AutoDiffIntegration.yml
index 39c3924af..90b02c5a6 100644
--- a/.github/workflows/AutoDiffIntegration.yml
+++ b/.github/workflows/AutoDiffIntegration.yml
@@ -14,7 +14,7 @@ concurrency:
 
 jobs:
   test:
-    name: Julia ${{ matrix.AD }} - ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.AD }} - ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -26,9 +26,6 @@ jobs:
           - ubuntu-latest
           - macOS-latest
           - windows-latest
-        arch:
-          - x64
-
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
diff --git a/.github/workflows/ForwardDiff.yml b/.github/workflows/ForwardDiff.yml
index 2100d7472..1008d371c 100644
--- a/.github/workflows/ForwardDiff.yml
+++ b/.github/workflows/ForwardDiff.yml
@@ -8,7 +8,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  Tests:
+  test:
     uses: ./.github/workflows/AutoDiffIntegration.yml
     with:
       AD: ForwardDiff
diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml
index 837f35200..a6548be3f 100644
--- a/.github/workflows/Tests.yml
+++ b/.github/workflows/Tests.yml
@@ -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,14 +25,11 @@ 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

From 7607a057c80a7fc990e8820d239c3fb62cc80f00 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Thu, 21 Aug 2025 12:03:21 -0400
Subject: [PATCH 15/21] fix bug in test selection

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

diff --git a/test/runtests.jl b/test/runtests.jl
index 9e5716c0f..9e88f936e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -54,7 +54,7 @@ end
 include("models/normal.jl")
 include("models/normallognormal.jl")
 
-if GROUP == "All" || GROUP == "General"
+if GROUP == "All" || GROUP == "General" || GROUP == "AD"
     if GROUP == "AD"
         include("general/ad.jl")
         include("general/optimize.jl")

From 4e08fc5fb9938555a6d262f78f94b9526b976978 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Thu, 21 Aug 2025 12:04:08 -0400
Subject: [PATCH 16/21] fix AD integration workflow

---
 .github/workflows/AutoDiffIntegration.yml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.github/workflows/AutoDiffIntegration.yml b/.github/workflows/AutoDiffIntegration.yml
index 90b02c5a6..3b02354a4 100644
--- a/.github/workflows/AutoDiffIntegration.yml
+++ b/.github/workflows/AutoDiffIntegration.yml
@@ -31,7 +31,6 @@ jobs:
       - 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

From 1dac1b3c9ad272894b1da1f655ca87d66ff591e9 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Thu, 21 Aug 2025 12:11:45 -0400
Subject: [PATCH 17/21] fix test only test mixed ad on the reversediff backends

---
 test/general/mixedad_logdensity.jl | 47 ++++++++++++++++--------------
 1 file changed, 25 insertions(+), 22 deletions(-)

diff --git a/test/general/mixedad_logdensity.jl b/test/general/mixedad_logdensity.jl
index cc7d0f89f..5e83ad4e9 100644
--- a/test/general/mixedad_logdensity.jl
+++ b/test/general/mixedad_logdensity.jl
@@ -31,27 +31,30 @@ 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 "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
+if AD isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoEnzyme,<:AutoMooncake}
+    @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

From f797d62e97a3fb72d4d4152a723f79ddb916ee75 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Thu, 21 Aug 2025 12:13:54 -0400
Subject: [PATCH 18/21] fix check for reverse-mode as well for mixed-ad

---
 test/general/mixedad_logdensity.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/general/mixedad_logdensity.jl b/test/general/mixedad_logdensity.jl
index 5e83ad4e9..bee48a397 100644
--- a/test/general/mixedad_logdensity.jl
+++ b/test/general/mixedad_logdensity.jl
@@ -31,7 +31,9 @@ function mixedad_test_fwd(x, prob)
     )/2
 end
 
-if AD isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoEnzyme,<:AutoMooncake}
+# 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)

From f048e133fc03fb9361a890b21bd90dc432e2089f Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Fri, 22 Aug 2025 12:47:12 -0400
Subject: [PATCH 19/21] fix remove redundant test

---
 test/algorithms/paramspacesgd/repgradelbo.jl | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/test/algorithms/paramspacesgd/repgradelbo.jl b/test/algorithms/paramspacesgd/repgradelbo.jl
index 48742db50..45f22c0ad 100644
--- a/test/algorithms/paramspacesgd/repgradelbo.jl
+++ b/test/algorithms/paramspacesgd/repgradelbo.jl
@@ -22,19 +22,6 @@
         end
     end
 
-    @testset "without mixed ad" begin
-        @testset for n_montecarlo in [1, 10]
-            alg = KLMinRepGradDescent(
-                AD;
-                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
-
     obj = RepGradELBO(10)
     rng = StableRNG(seed)
     elbo_ref = estimate_objective(rng, obj, q0, model; n_samples=10^4)

From 7b8dad31d9b7f794214dd27cd0735bd97ecf8378 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Fri, 22 Aug 2025 12:48:38 -0400
Subject: [PATCH 20/21] refactor reorganize tests

---
 .github/workflows/Enzyme.yml      |  2 +-
 .github/workflows/ForwardDiff.yml |  2 +-
 .github/workflows/ReverseDiff.yml | 14 ++++++++++++++
 .github/workflows/Tests.yml       |  2 +-
 test/runtests.jl                  | 20 ++++++++++----------
 5 files changed, 27 insertions(+), 13 deletions(-)
 create mode 100644 .github/workflows/ReverseDiff.yml

diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml
index cf6555f30..eb97527c2 100644
--- a/.github/workflows/Enzyme.yml
+++ b/.github/workflows/Enzyme.yml
@@ -1,4 +1,4 @@
-name: Enzyme Integration 
+name: Enzyme Integration Tests
 on:
   push:
     branches:
diff --git a/.github/workflows/ForwardDiff.yml b/.github/workflows/ForwardDiff.yml
index 1008d371c..830ae9a4e 100644
--- a/.github/workflows/ForwardDiff.yml
+++ b/.github/workflows/ForwardDiff.yml
@@ -1,4 +1,4 @@
-name: ForwardDiff Integration 
+name: ForwardDiff Integration Tests
 on:
   push:
     branches:
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/Tests.yml b/.github/workflows/Tests.yml
index a6548be3f..51a6bfeb0 100644
--- a/.github/workflows/Tests.yml
+++ b/.github/workflows/Tests.yml
@@ -1,4 +1,4 @@
-name: All Tests with ReverseDiff Integration
+name: General Tests
 on:
   push:
     branches:
diff --git a/test/runtests.jl b/test/runtests.jl
index 9e88f936e..18e503dad 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -54,24 +54,24 @@ end
 include("models/normal.jl")
 include("models/normallognormal.jl")
 
-if GROUP == "All" || GROUP == "General" || GROUP == "AD"
-    if GROUP == "AD"
-        include("general/ad.jl")
-        include("general/optimize.jl")
-        include("general/proximal_location_scale_entropy.jl")
-        include("general/mixedad_logdensity.jl")
-    end
+
+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")
-end
 
-if GROUP == "All" || GROUP == "General" || GROUP == "Families"
     include("families/location_scale.jl")
     include("families/location_scale_low_rank.jl")
 end
 
-if GROUP == "All" || GROUP == "ParamSpaceSGD" || GROUP == "AD"
+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")

From 592820121a0960bc0589e8a9908699b1f2b4a773 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <msca8h@naver.com>
Date: Fri, 22 Aug 2025 12:51:35 -0400
Subject: [PATCH 21/21] run formatter

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 test/runtests.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 18e503dad..3e33b7405 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -54,7 +54,6 @@ end
 include("models/normal.jl")
 include("models/normallognormal.jl")
 
-
 if GROUP == "All" || GROUP == "GENERAL"
     # Tests that do not need to check correct integration with AD backends
     include("general/optimize.jl")