From 4c547f3ea7d4ef7cf457015597c5e240188d05cf Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Tue, 19 Aug 2025 16:33:16 -0400 Subject: [PATCH 1/3] drop tests with DistributionsAD --- test/Project.toml | 2 - .../repgradelbo_distributionsad.jl | 75 ------------------- .../scoregradelbo_distributionsad.jl | 69 ----------------- test/runtests.jl | 10 +-- 4 files changed, 2 insertions(+), 154 deletions(-) delete mode 100644 test/algorithms/paramspacesgd/repgradelbo_distributionsad.jl delete mode 100644 test/algorithms/paramspacesgd/scoregradelbo_distributionsad.jl diff --git a/test/Project.toml b/test/Project.toml index b626b1218..db7d9068a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,7 +4,6 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -30,7 +29,6 @@ Bijectors = "0.13, 0.14, 0.15" DiffResults = "1" DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" -DistributionsAD = "0.6.45" Enzyme = "0.13, 0.14, 0.15" FillArrays = "1.6.1" ForwardDiff = "0.10.36, 1" diff --git a/test/algorithms/paramspacesgd/repgradelbo_distributionsad.jl b/test/algorithms/paramspacesgd/repgradelbo_distributionsad.jl deleted file mode 100644 index 0180a2eb3..000000000 --- a/test/algorithms/paramspacesgd/repgradelbo_distributionsad.jl +++ /dev/null @@ -1,75 +0,0 @@ - -AD_repgradelbo_distributionsad = if TEST_GROUP == "Enzyme" - Dict( - :Enzyme => AutoEnzyme(; - mode=Enzyme.set_runtime_activity(Enzyme.Reverse), - function_annotation=Enzyme.Const, - ), - ) -else - Dict( - #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment - :Zygote => AutoZygote(), - :Mooncake => AutoMooncake(; config=Mooncake.Config()), - ) -end - -@testset "inference RepGradELBO DistributionsAD" begin - @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in - [Float64, Float32], - (modelname, modelconstr) in Dict(:Normal => normal_meanfield), - (objname, objective) in Dict( - :RepGradELBOClosedFormEntropy => RepGradELBO(10), - :RepGradELBOStickingTheLanding => - RepGradELBO(10; entropy=StickingTheLandingEntropy()), - ), - (adbackname, adtype) in AD_repgradelbo_distributionsad - - seed = (0x38bef07cf9cc549d) - rng = StableRNG(seed) - - modelstats = modelconstr(rng, realtype) - (; model, μ_true, L_true, n_dims, strong_convexity, is_meanfield) = modelstats - - μ0 = zeros(realtype, n_dims) - L0 = Diagonal(ones(realtype, n_dims)) - q0 = TuringDiagMvNormal(μ0, diag(L0)) - - T = 1000 - η = 1e-3 - alg = KLMinRepGradDescent(adtype; operator=IdentityOperator(), optimizer=Descent(η)) - - # For small enough η, the error of SGD, Δλ, is bounded as - # Δλ ≤ ρ^T Δλ0 + O(η), - # where ρ = 1 - ημ, μ is the strong convexity constant. - contraction_rate = 1 - η * strong_convexity - - @testset "convergence" begin - Δλ0 = sum(abs2, μ0 - μ_true) + sum(abs2, L0 - L_true) - - q_avg, stats, _ = optimize(rng, alg, T, model, q0; show_progress=PROGRESS) - - μ = mean(q_avg) - L = sqrt(cov(q_avg)) - Δλ = sum(abs2, μ - μ_true) + sum(abs2, L - L_true) - - @test Δλ ≤ contraction_rate^(T / 2) * Δλ0 - @test eltype(μ) == eltype(μ_true) - @test eltype(L) == eltype(L_true) - end - - @testset "determinism" begin - rng = StableRNG(seed) - q_avg, stats, _ = optimize(rng, alg, T, model, q0; show_progress=PROGRESS) - μ = mean(q_avg) - L = sqrt(cov(q_avg)) - - rng_repl = StableRNG(seed) - q_avg, stats, _ = optimize(rng_repl, alg, T, model, q0; show_progress=PROGRESS) - μ_repl = mean(q_avg) - L_repl = sqrt(cov(q_avg)) - @test μ == μ_repl - @test L == L_repl - end - end -end diff --git a/test/algorithms/paramspacesgd/scoregradelbo_distributionsad.jl b/test/algorithms/paramspacesgd/scoregradelbo_distributionsad.jl deleted file mode 100644 index a9a505b26..000000000 --- a/test/algorithms/paramspacesgd/scoregradelbo_distributionsad.jl +++ /dev/null @@ -1,69 +0,0 @@ -AD_scoregradelbo_distributionsad = 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 DistributionsAD" begin - @testset "$(modelname) $(realtype) $(adbackname)" for realtype in [Float64, Float32], - (modelname, modelconstr) in Dict(:Normal => normal_meanfield), - (adbackname, adtype) in AD_scoregradelbo_distributionsad - - seed = (0x38bef07cf9cc549d) - rng = StableRNG(seed) - - modelstats = modelconstr(rng, realtype) - (; model, μ_true, L_true, n_dims, strong_convexity, is_meanfield) = modelstats - - T = 1000 - η = 1e-4 - opt = Optimisers.Descent(η) - alg = KLMinScoreGradDescent(adtype; n_samples=10, optimizer=opt) - - # For small enough η, the error of SGD, Δλ, is bounded as - # Δλ ≤ ρ^T Δλ0 + O(η), - # where ρ = 1 - ημ, μ is the strong convexity constant. - contraction_rate = 1 - η * strong_convexity - - μ0 = zeros(realtype, n_dims) - L0 = Diagonal(ones(realtype, n_dims)) - q0 = TuringDiagMvNormal(μ0, diag(L0)) - - @testset "convergence" begin - Δλ0 = sum(abs2, μ0 - μ_true) + sum(abs2, L0 - L_true) - q_avg, stats, _ = optimize(rng, alg, T, model, q0; show_progress=PROGRESS) - - μ = mean(q_avg) - L = sqrt(cov(q_avg)) - Δλ = sum(abs2, μ - μ_true) + sum(abs2, L - L_true) - - @test Δλ ≤ contraction_rate^(T / 2) * Δλ0 - @test eltype(μ) == eltype(μ_true) - @test eltype(L) == eltype(L_true) - end - - @testset "determinism" begin - rng = StableRNG(seed) - q_avg, stats, _ = optimize(rng, alg, T, model, q0; show_progress=PROGRESS) - μ = mean(q_avg) - L = sqrt(cov(q_avg)) - - rng_repl = StableRNG(seed) - q_avg, stats, _ = optimize(rng_repl, alg, T, model, q0; show_progress=PROGRESS) - μ_repl = mean(q_avg) - L_repl = sqrt(cov(q_avg)) - @test μ ≈ μ_repl rtol = 1e-5 - @test L ≈ L_repl rtol = 1e-5 - end - end -end diff --git a/test/runtests.jl b/test/runtests.jl index aea922fff..d0cda83b7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,19 +16,15 @@ using Random, StableRNGs using Statistics using StatsBase -using Functors -using DistributionsAD -@functor TuringDiagMvNormal - 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") -if TEST_GROUP == "Enzyme" +if GROUP == "Enzyme" using Enzyme end @@ -67,12 +63,10 @@ end if TEST_GROUP == "All" || TEST_GROUP == "ParamSpaceSGD" || TEST_GROUP == "Enzyme" include("algorithms/paramspacesgd/repgradelbo.jl") include("algorithms/paramspacesgd/scoregradelbo.jl") - include("algorithms/paramspacesgd/repgradelbo_distributionsad.jl") include("algorithms/paramspacesgd/repgradelbo_locationscale.jl") include("algorithms/paramspacesgd/repgradelbo_locationscale_bijectors.jl") include("algorithms/paramspacesgd/repgradelbo_proximal_locationscale.jl") include("algorithms/paramspacesgd/repgradelbo_proximal_locationscale_bijectors.jl") - include("algorithms/paramspacesgd/scoregradelbo_distributionsad.jl") include("algorithms/paramspacesgd/scoregradelbo_locationscale.jl") include("algorithms/paramspacesgd/scoregradelbo_locationscale_bijectors.jl") end From 435a3a4ff952c8741079c911997aa1c54615b8d9 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Tue, 19 Aug 2025 16:39:38 -0400 Subject: [PATCH 2/3] fix revert wrong change --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index d0cda83b7..99ef5def1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,7 +22,7 @@ using ForwardDiff, ReverseDiff, Zygote, Mooncake using AdvancedVI const PROGRESS = haskey(ENV, "PROGRESS") -const GROUP = get(ENV, "GROUP", "All") +const TEST_GROUP = get(ENV, "TEST_GROUP", "All") if GROUP == "Enzyme" using Enzyme From c2416f05a3d496bed28a91677bd67e273db10203 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Tue, 19 Aug 2025 16:43:47 -0400 Subject: [PATCH 3/3] fix wrong variable name --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 99ef5def1..e72e00af8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,7 +24,7 @@ using AdvancedVI const PROGRESS = haskey(ENV, "PROGRESS") const TEST_GROUP = get(ENV, "TEST_GROUP", "All") -if GROUP == "Enzyme" +if TEST_GROUP == "Enzyme" using Enzyme end