Split AD tests into workflows

.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
+

.github/workflows/Enzyme.yml

@@ -1,40 +1,14 @@
-name: Enzyme
+name: Enzyme Integration Tests
 on:
   push:
     branches:
       - main
     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

.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

.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

.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

.github/workflows/CI.yml → .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:

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,19 +45,17 @@ 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)))
         )
         params, re = Optimisers.destructure(q_true)
         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

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)))

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)

test/algorithms/paramspacesgd/repgradelbo_proximal_locationscale.jl

@@ -1,31 +1,13 @@
-
-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(
             :RepGradELBOClosedFormEntropy =>
                 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(η),