From b23428b9704d4bcf3c55a9233f71868072235879 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Fri, 21 Jun 2024 01:29:06 -0400 Subject: [PATCH 01/15] fix enzyme to match new interface, enable enzyme tests --- ext/AdvancedVIEnzymeExt.jl | 42 +++++++++++++++---- test/Project.toml | 3 ++ test/inference/repgradelbo_locationscale.jl | 2 +- .../repgradelbo_locationscale_bijectors.jl | 4 +- test/interface/ad.jl | 21 +++++++++- test/runtests.jl | 2 +- 6 files changed, 62 insertions(+), 12 deletions(-) diff --git a/ext/AdvancedVIEnzymeExt.jl b/ext/AdvancedVIEnzymeExt.jl index 8333299f..a334ca64 100644 --- a/ext/AdvancedVIEnzymeExt.jl +++ b/ext/AdvancedVIEnzymeExt.jl @@ -11,15 +11,43 @@ else using ..AdvancedVI: ADTypes, DiffResults end -# Enzyme doesn't support f::Bijectors (see https://github.com/EnzymeAD/Enzyme.jl/issues/916) function AdvancedVI.value_and_gradient!( - ad::ADTypes.AutoEnzyme, f, θ::AbstractVector{T}, out::DiffResults.MutableDiffResult -) where {T<:Real} - y = f(θ) + ::ADTypes.AutoEnzyme, + f, + x ::AbstractVector{<:Real}, + out::DiffResults.MutableDiffResult +) + Enzyme.API.runtimeActivity!(true) + ∇x = DiffResults.gradient(out) + fill!(∇x, zero(eltype(∇x))) + _, y = Enzyme.autodiff( + Enzyme.ReverseWithPrimal, + f, + Enzyme.Active, + Enzyme.Duplicated(x, ∇x) + ) + DiffResults.value!(out, y) + return out +end + +function AdvancedVI.value_and_gradient!( + ::ADTypes.AutoEnzyme, + f, + x ::AbstractVector{<:Real}, + aux, + out ::DiffResults.MutableDiffResult +) + Enzyme.API.runtimeActivity!(true) + ∇x = DiffResults.gradient(out) + fill!(∇x, zero(eltype(∇x))) + _, y = Enzyme.autodiff( + Enzyme.ReverseWithPrimal, + f, + Enzyme.Active, + Enzyme.Duplicated(x, ∇x), + Enzyme.Const(aux) + ) DiffResults.value!(out, y) - ∇θ = DiffResults.gradient(out) - fill!(∇θ, zero(T)) - Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(θ, ∇θ)) return out end diff --git a/test/Project.toml b/test/Project.toml index a0dba17f..07997916 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" 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" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" @@ -24,8 +25,10 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] ADTypes = "0.2.1, 1" Bijectors = "0.13" +DiffResults = "1.0" Distributions = "0.25.100" DistributionsAD = "0.6.45" +Enzyme = "0.12" FillArrays = "1.6.1" ForwardDiff = "0.10.36" Functors = "0.4.5" diff --git a/test/inference/repgradelbo_locationscale.jl b/test/inference/repgradelbo_locationscale.jl index 73846675..cb255226 100644 --- a/test/inference/repgradelbo_locationscale.jl +++ b/test/inference/repgradelbo_locationscale.jl @@ -15,7 +15,7 @@ :ForwarDiff => AutoForwardDiff(), :ReverseDiff => AutoReverseDiff(), :Zygote => AutoZygote(), - #:Enzyme => AutoEnzyme(), + :Enzyme => AutoEnzyme(), ) seed = (0x38bef07cf9cc549d) diff --git a/test/inference/repgradelbo_locationscale_bijectors.jl b/test/inference/repgradelbo_locationscale_bijectors.jl index 1f62df0f..d56af333 100644 --- a/test/inference/repgradelbo_locationscale_bijectors.jl +++ b/test/inference/repgradelbo_locationscale_bijectors.jl @@ -13,8 +13,8 @@ (adbackname, adtype) in Dict( :ForwarDiff => AutoForwardDiff(), :ReverseDiff => AutoReverseDiff(), - #:Zygote => AutoZygote(), - #:Enzyme => AutoEnzyme(), + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), ) seed = (0x38bef07cf9cc549d) diff --git a/test/interface/ad.jl b/test/interface/ad.jl index be4ca34e..faa1f01c 100644 --- a/test/interface/ad.jl +++ b/test/interface/ad.jl @@ -6,7 +6,7 @@ using Test :ForwardDiff => AutoForwardDiff(), :ReverseDiff => AutoReverseDiff(), :Zygote => AutoZygote(), - # :Enzyme => AutoEnzyme() # Currently not tested against + :Enzyme => AutoEnzyme(), ) D = 10 A = randn(D, D) @@ -19,4 +19,23 @@ using Test @test ∇ ≈ (A + A')*λ/2 @test f ≈ λ'*A*λ / 2 end + + @testset "$(adname) with auxiliary input" for (adname, adsymbol) ∈ Dict( + :ForwardDiff => AutoForwardDiff(), + :ReverseDiff => AutoReverseDiff(), + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), + ) + 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!(adsymbol, f, λ, (b=b,), grad_buf) + ∇ = DiffResults.gradient(grad_buf) + f = DiffResults.value(grad_buf) + @test ∇ ≈ (A + A')*λ/2 + b + @test f ≈ λ'*A*λ / 2 + dot(b, λ) + end end diff --git a/test/runtests.jl b/test/runtests.jl index 3bd13144..29850b65 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -18,7 +18,7 @@ using DistributionsAD using LogDensityProblems using Optimisers using ADTypes -using ForwardDiff, ReverseDiff, Zygote +using ForwardDiff, ReverseDiff, Zygote, Enzyme using AdvancedVI From ed38340e49b867991b59a855703244cffeee93ab Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Mon, 1 Jul 2024 16:06:25 -0400 Subject: [PATCH 02/15] fix type instability tighten Enzyme compat --- Project.toml | 2 +- src/objectives/elbo/repgradelbo.jl | 2 +- test/Project.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 9d30e740..a35a5bd4 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ ChainRulesCore = "1.16" DiffResults = "1" Distributions = "0.25.87" DocStringExtensions = "0.8, 0.9" -Enzyme = "0.12" +Enzyme = "0.12.17" FillArrays = "1.3" ForwardDiff = "0.10.36" Functors = "0.4" diff --git a/src/objectives/elbo/repgradelbo.jl b/src/objectives/elbo/repgradelbo.jl index 27a937e8..227887ce 100644 --- a/src/objectives/elbo/repgradelbo.jl +++ b/src/objectives/elbo/repgradelbo.jl @@ -99,7 +99,7 @@ estimate_objective(obj::RepGradELBO, q, prob; n_samples::Int = obj.n_samples) = function estimate_repgradelbo_ad_forward(params′, aux) @unpack rng, obj, problem, restructure, q_stop = aux - q = restructure(params′) + q = restructure(params′)::typeof(restructure.model) samples, entropy = reparam_with_entropy(rng, q, q_stop, obj.n_samples, obj.entropy) energy = estimate_energy_with_samples(problem, samples) elbo = energy + entropy diff --git a/test/Project.toml b/test/Project.toml index 07997916..b79589d1 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,7 +28,7 @@ Bijectors = "0.13" DiffResults = "1.0" Distributions = "0.25.100" DistributionsAD = "0.6.45" -Enzyme = "0.12" +Enzyme = "0.12.17" FillArrays = "1.6.1" ForwardDiff = "0.10.36" Functors = "0.4.5" From 2af7695bde58b0e8df3bebbea5c3b577cd4d2fe5 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sun, 14 Jul 2024 22:50:11 -0400 Subject: [PATCH 03/15] add indirection to enforce type stability of `restructure` --- ext/AdvancedVIEnzymeExt.jl | 4 ++++ src/AdvancedVI.jl | 17 +++++++++-------- src/families/location_scale.jl | 4 ++-- src/objectives/elbo/repgradelbo.jl | 13 ++++++++++--- test/inference/repgradelbo_distributionsad.jl | 2 +- test/interface/repgradelbo.jl | 5 +++-- 6 files changed, 29 insertions(+), 16 deletions(-) diff --git a/ext/AdvancedVIEnzymeExt.jl b/ext/AdvancedVIEnzymeExt.jl index a334ca64..a0202553 100644 --- a/ext/AdvancedVIEnzymeExt.jl +++ b/ext/AdvancedVIEnzymeExt.jl @@ -11,6 +11,10 @@ else using ..AdvancedVI: ADTypes, DiffResults end +AdvancedVI.restructure_ad_forward( + ::ADTypes.AutoEnzyme, restructure, params +) = restructure(params)::typeof(restructure.model) + function AdvancedVI.value_and_gradient!( ::ADTypes.AutoEnzyme, f, diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 7a09030b..03fcff5a 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -40,19 +40,20 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif """ function value_and_gradient! end + +# derivatives """ - stop_gradient(x) + restructure_ad_forward(adtype, restructure, params) -Stop the gradient from propagating to `x` if the selected ad backend supports it. -Otherwise, it is equivalent to `identity`. +Apply `restructure` to `params`. +This is an indirection for handling the type stability of `restructure`, as some AD backends require strict type stability in the AD path. # Arguments -- `x`: Input - -# Returns -- `x`: Same value as the input. +- `ad::ADTypes.AbstractADType`: Automatic differentiation backend. +- `restructure`: Callable for restructuring the varitional distribution from `params`. +- `params`: Variational Parameters. """ -function stop_gradient end +restructure_ad_forward(::ADTypes.AbstractADType, restructure, params) = restructure(params) # Update for gradient descent step """ diff --git a/src/families/location_scale.jl b/src/families/location_scale.jl index 66ea5cdb..f86257db 100644 --- a/src/families/location_scale.jl +++ b/src/families/location_scale.jl @@ -39,14 +39,14 @@ Functors.@functor MvLocationScale (location, scale) # is very inefficient. # begin struct RestructureMeanField{S <: Diagonal, D, L} - q::MvLocationScale{S, D, L} + model::MvLocationScale{S, D, L} end function (re::RestructureMeanField)(flat::AbstractVector) n_dims = div(length(flat), 2) location = first(flat, n_dims) scale = Diagonal(last(flat, n_dims)) - MvLocationScale(location, scale, re.q.dist, re.q.scale_eps) + MvLocationScale(location, scale, re.model.dist, re.model.scale_eps) end function Optimisers.destructure( diff --git a/src/objectives/elbo/repgradelbo.jl b/src/objectives/elbo/repgradelbo.jl index 227887ce..2daf880b 100644 --- a/src/objectives/elbo/repgradelbo.jl +++ b/src/objectives/elbo/repgradelbo.jl @@ -98,8 +98,8 @@ estimate_objective(obj::RepGradELBO, q, prob; n_samples::Int = obj.n_samples) = estimate_objective(Random.default_rng(), obj, q, prob; n_samples) function estimate_repgradelbo_ad_forward(params′, aux) - @unpack rng, obj, problem, restructure, q_stop = aux - q = restructure(params′)::typeof(restructure.model) + @unpack rng, obj, problem, adtype, restructure, q_stop = aux + q = restructure_ad_forward(adtype, restructure, params′) samples, entropy = reparam_with_entropy(rng, q, q_stop, obj.n_samples, obj.entropy) energy = estimate_energy_with_samples(problem, samples) elbo = energy + entropy @@ -117,7 +117,14 @@ function estimate_gradient!( state, ) q_stop = restructure(params) - aux = (rng=rng, obj=obj, problem=prob, restructure=restructure, q_stop=q_stop) + aux = ( + rng = rng, + adtype = adtype, + obj = obj, + problem = prob, + restructure = restructure, + q_stop = q_stop + ) value_and_gradient!( adtype, estimate_repgradelbo_ad_forward, params, aux, out ) diff --git a/test/inference/repgradelbo_distributionsad.jl b/test/inference/repgradelbo_distributionsad.jl index 38dbc6e3..d443e569 100644 --- a/test/inference/repgradelbo_distributionsad.jl +++ b/test/inference/repgradelbo_distributionsad.jl @@ -14,7 +14,7 @@ :ForwarDiff => AutoForwardDiff(), #:ReverseDiff => AutoReverseDiff(), :Zygote => AutoZygote(), - #:Enzyme => AutoEnzyme(), + :Enzyme => AutoEnzyme(), ) seed = (0x38bef07cf9cc549d) diff --git a/test/interface/repgradelbo.jl b/test/interface/repgradelbo.jl index ac9bfeca..579aba78 100644 --- a/test/interface/repgradelbo.jl +++ b/test/interface/repgradelbo.jl @@ -37,7 +37,8 @@ end @testset for ad in [ ADTypes.AutoForwardDiff(), ADTypes.AutoReverseDiff(), - ADTypes.AutoZygote() + ADTypes.AutoZygote(), + ADTypes.AutoEnzyme() ] q_true = MeanFieldGaussian( Vector{eltype(μ_true)}(μ_true), @@ -47,7 +48,7 @@ end obj = RepGradELBO(10; entropy=StickingTheLandingEntropy()) out = DiffResults.DiffResult(zero(eltype(params)), similar(params)) - aux = (rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true) + aux = (rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true, adtype=ad) AdvancedVI.value_and_gradient!( ad, AdvancedVI.estimate_repgradelbo_ad_forward, params, aux, out ) From a5be89dc09d16af408b7b2611c5238ff3299e322 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Wed, 17 Jul 2024 00:47:12 -0400 Subject: [PATCH 04/15] minor formatting changes, fix DistributionsAD inference test --- test/inference/repgradelbo_distributionsad.jl | 8 ++++---- test/models/normal.jl | 2 -- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/test/inference/repgradelbo_distributionsad.jl b/test/inference/repgradelbo_distributionsad.jl index d443e569..c139facc 100644 --- a/test/inference/repgradelbo_distributionsad.jl +++ b/test/inference/repgradelbo_distributionsad.jl @@ -3,7 +3,7 @@ @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype ∈ [Float64, Float32], (modelname, modelconstr) ∈ Dict( - :Normal=> normal_meanfield, + :Normal => normal_meanfield, ), n_montecarlo in [1, 10], (objname, objective) in Dict( @@ -12,7 +12,7 @@ ), (adbackname, adtype) ∈ Dict( :ForwarDiff => AutoForwardDiff(), - #:ReverseDiff => AutoReverseDiff(), + #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment :Zygote => AutoZygote(), :Enzyme => AutoEnzyme(), ) @@ -32,8 +32,8 @@ # where ρ = 1 - ημ, μ is the strong convexity constant. contraction_rate = 1 - η*strong_convexity - μ0 = Zeros(realtype, n_dims) - L0 = Diagonal(Ones(realtype, n_dims)) + μ0 = zeros(realtype, n_dims) + L0 = Diagonal(ones(realtype, n_dims)) q0 = TuringDiagMvNormal(μ0, diag(L0)) @testset "convergence" begin diff --git a/test/models/normal.jl b/test/models/normal.jl index 5c3e22e8..59cf0043 100644 --- a/test/models/normal.jl +++ b/test/models/normal.jl @@ -35,9 +35,7 @@ function normal_meanfield(rng::Random.AbstractRNG, realtype::Type) σ0 = realtype(0.3) μ = Fill(realtype(5), n_dims) - #randn(rng, realtype, n_dims) σ = Fill(σ0, n_dims) - #log.(exp.(randn(rng, realtype, n_dims)) .+ 1) model = TestNormal(μ, Diagonal(σ.^2)) From 1727f5f85e3c6c105092a0a455582e3ea106d92f Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Wed, 17 Jul 2024 00:49:15 -0400 Subject: [PATCH 05/15] tighten compat bound for Enzyme --- Project.toml | 2 +- test/Project.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index a35a5bd4..40c6dc4f 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ ChainRulesCore = "1.16" DiffResults = "1" Distributions = "0.25.87" DocStringExtensions = "0.8, 0.9" -Enzyme = "0.12.17" +Enzyme = "0.12.23" FillArrays = "1.3" ForwardDiff = "0.10.36" Functors = "0.4" diff --git a/test/Project.toml b/test/Project.toml index b79589d1..91c70090 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,7 +28,7 @@ Bijectors = "0.13" DiffResults = "1.0" Distributions = "0.25.100" DistributionsAD = "0.6.45" -Enzyme = "0.12.17" +Enzyme = "0.12.23" FillArrays = "1.6.1" ForwardDiff = "0.10.36" Functors = "0.4.5" From 69a475480f563ec84fc0e395cddcf75d4b082efa Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 3 Aug 2024 00:16:30 -0400 Subject: [PATCH 06/15] fix remove trailing whitespace in docs --- docs/make.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index 7ae3bc62..d23b9f24 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -15,8 +15,9 @@ makedocs(; "General Usage" => "general.md", "Examples" => "examples.md", "ELBO Maximization" => [ - "Overview" => "elbo/overview.md", - "Reparameterization Gradient Estimator" => "elbo/repgradelbo.md", + "Overview" => "elbo/overview.md", + "Reparameterization Gradient Estimator" => "elbo/repgradelbo.md", + "Sample Average Approximation" => "elbo/saa.md", "Location-Scale Variational Family" => "locscale.md", ]], ) From ef43c6c1174292bd1d6e34832019b279ea5a5444 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sun, 4 Aug 2024 23:48:08 -0400 Subject: [PATCH 07/15] tighten compat bound for enzyme --- test/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index 91c70090..7b612667 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,7 +28,7 @@ Bijectors = "0.13" DiffResults = "1.0" Distributions = "0.25.100" DistributionsAD = "0.6.45" -Enzyme = "0.12.23" +Enzyme = "0.12.27" FillArrays = "1.6.1" ForwardDiff = "0.10.36" Functors = "0.4.5" From 156398c35c56aaf02cac36852a16fc0756e3f366 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Wed, 7 Aug 2024 00:59:09 -0400 Subject: [PATCH 08/15] remove unused comment --- src/AdvancedVI.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 03fcff5a..a7f0b6aa 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -40,8 +40,6 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif """ function value_and_gradient! end - -# derivatives """ restructure_ad_forward(adtype, restructure, params) From d821554e0b9ba9428954e9e44794b080f5105482 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 10 Aug 2024 01:02:16 -0400 Subject: [PATCH 09/15] compat tighten julia version --- .github/workflows/CI.yml | 3 +-- Project.toml | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 0302271a..b19874b3 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -19,8 +19,7 @@ jobs: fail-fast: false matrix: version: - - '1' - - '1.6' + - '1.7' os: - ubuntu-latest - macOS-latest diff --git a/Project.toml b/Project.toml index 7a6a75d4..a5bc618b 100644 --- a/Project.toml +++ b/Project.toml @@ -56,7 +56,7 @@ ReverseDiff = "1.15.1" SimpleUnPack = "1.1.0" StatsBase = "0.32, 0.33, 0.34" Zygote = "0.6.63" -julia = "1.6" +julia = "1.7" [extras] Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" From 36d22cfded28191e69542804651c40df89afad4e Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 10 Aug 2024 01:04:52 -0400 Subject: [PATCH 10/15] fix formatting --- docs/make.jl | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index d23b9f24..b1785ef2 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,24 +2,23 @@ using AdvancedVI using Documenter -DocMeta.setdocmeta!( - AdvancedVI, :DocTestSetup, :(using AdvancedVI); recursive=true -) +DocMeta.setdocmeta!(AdvancedVI, :DocTestSetup, :(using AdvancedVI); recursive=true) makedocs(; - modules = [AdvancedVI], - sitename = "AdvancedVI.jl", - repo = "https://github.com/TuringLang/AdvancedVI.jl/blob/{commit}{path}#{line}", - format = Documenter.HTML(; prettyurls = get(ENV, "CI", nothing) == "true"), - pages = ["AdvancedVI" => "index.md", - "General Usage" => "general.md", - "Examples" => "examples.md", - "ELBO Maximization" => [ - "Overview" => "elbo/overview.md", - "Reparameterization Gradient Estimator" => "elbo/repgradelbo.md", - "Sample Average Approximation" => "elbo/saa.md", - "Location-Scale Variational Family" => "locscale.md", - ]], + modules=[AdvancedVI], + sitename="AdvancedVI.jl", + repo="https://github.com/TuringLang/AdvancedVI.jl/blob/{commit}{path}#{line}", + format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true"), + pages=[ + "AdvancedVI" => "index.md", + "General Usage" => "general.md", + "Examples" => "examples.md", + "ELBO Maximization" => [ + "Overview" => "elbo/overview.md", + "Reparameterization Gradient Estimator" => "elbo/repgradelbo.md", + "Location-Scale Variational Family" => "locscale.md", + ], + ], ) deploydocs(; repo="github.com/TuringLang/AdvancedVI.jl", push_preview=true) From c7f60a5ef2ba1cde324aa865a06ba20fc38324d5 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 10 Aug 2024 01:05:19 -0400 Subject: [PATCH 11/15] fix formatting --- test/inference/repgradelbo_locationscale_bijectors.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/inference/repgradelbo_locationscale_bijectors.jl b/test/inference/repgradelbo_locationscale_bijectors.jl index 83bc858f..f9506c17 100644 --- a/test/inference/repgradelbo_locationscale_bijectors.jl +++ b/test/inference/repgradelbo_locationscale_bijectors.jl @@ -13,8 +13,8 @@ (adbackname, adtype) in Dict( :ForwarDiff => AutoForwardDiff(), :ReverseDiff => AutoReverseDiff(), - :Zygote => AutoZygote(), - :Enzyme => AutoEnzyme(), + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), ) seed = (0x38bef07cf9cc549d) From b0cd0b5c9aa0e6515be1d44325b03b7bfc5e8634 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 10 Aug 2024 02:21:32 -0400 Subject: [PATCH 12/15] try distable `DistributionsAD` --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index a827dbf9..2f1d775b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -47,7 +47,7 @@ end const PROGRESS = haskey(ENV, "PROGRESS") if GROUP == "All" || GROUP == "Inference" - include("inference/repgradelbo_distributionsad.jl") + #include("inference/repgradelbo_distributionsad.jl") include("inference/repgradelbo_locationscale.jl") include("inference/repgradelbo_locationscale_bijectors.jl") end From 72ed3fb811477f1dd7e1db756f8aacd7de63cc35 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 10 Aug 2024 23:12:54 -0400 Subject: [PATCH 13/15] fix tests enable Enzyme inference tests only on 1.10 --- .github/workflows/CI.yml | 3 ++- test/inference/repgradelbo_distributionsad.jl | 22 ++++++++++++++----- test/inference/repgradelbo_locationscale.jl | 22 ++++++++++++++----- .../repgradelbo_locationscale_bijectors.jl | 22 ++++++++++++++----- test/runtests.jl | 2 +- 5 files changed, 51 insertions(+), 20 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index b19874b3..7b2f199e 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -19,7 +19,8 @@ jobs: fail-fast: false matrix: version: - - '1.7' + - '1.6' + - '1.10' os: - ubuntu-latest - macOS-latest diff --git a/test/inference/repgradelbo_distributionsad.jl b/test/inference/repgradelbo_distributionsad.jl index e743f437..bdce1293 100644 --- a/test/inference/repgradelbo_distributionsad.jl +++ b/test/inference/repgradelbo_distributionsad.jl @@ -1,4 +1,19 @@ +AD_distributionsad = if VERSION >= v"1.10" + Dict( + :ForwarDiff => AutoForwardDiff(), + #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), + ) +else + Dict( + :ForwarDiff => AutoForwardDiff(), + #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment + :Zygote => AutoZygote(), + ) +end + @testset "inference RepGradELBO DistributionsAD" begin @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in [Float64, Float32], @@ -9,12 +24,7 @@ :RepGradELBOStickingTheLanding => RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()), ), - (adbackname, adtype) in Dict( - :ForwarDiff => AutoForwardDiff(), - #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment - :Zygote => AutoZygote(), - :Enzyme => AutoEnzyme(), - ) + (adbackname, adtype) in AD_distributionsad seed = (0x38bef07cf9cc549d) rng = StableRNG(seed) diff --git a/test/inference/repgradelbo_locationscale.jl b/test/inference/repgradelbo_locationscale.jl index 2fcaa421..bf94b270 100644 --- a/test/inference/repgradelbo_locationscale.jl +++ b/test/inference/repgradelbo_locationscale.jl @@ -1,4 +1,19 @@ +AD_locationscale = if VERSION >= v"1.10" + Dict( + :ForwarDiff => AutoForwardDiff(), + :ReverseDiff => AutoReverseDiff(), + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), + ) +else + Dict( + :ForwarDiff => AutoForwardDiff(), + :ReverseDiff => AutoReverseDiff(), + :Zygote => AutoZygote(), + ) +end + @testset "inference RepGradELBO VILocationScale" begin @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in [Float64, Float32], @@ -10,12 +25,7 @@ :RepGradELBOStickingTheLanding => RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()), ), - (adbackname, adtype) in Dict( - :ForwarDiff => AutoForwardDiff(), - :ReverseDiff => AutoReverseDiff(), - :Zygote => AutoZygote(), - :Enzyme => AutoEnzyme(), - ) + (adbackname, adtype) in AD_locationscale seed = (0x38bef07cf9cc549d) rng = StableRNG(seed) diff --git a/test/inference/repgradelbo_locationscale_bijectors.jl b/test/inference/repgradelbo_locationscale_bijectors.jl index f9506c17..80150e9c 100644 --- a/test/inference/repgradelbo_locationscale_bijectors.jl +++ b/test/inference/repgradelbo_locationscale_bijectors.jl @@ -1,4 +1,19 @@ +AD_locationscale_bijectors = if VERSION >= v"1.10" + Dict( + :ForwarDiff => AutoForwardDiff(), + :ReverseDiff => AutoReverseDiff(), + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), + ) +else + Dict( + :ForwarDiff => AutoForwardDiff(), + :ReverseDiff => AutoReverseDiff(), + :Zygote => AutoZygote(), + ) +end + @testset "inference RepGradELBO VILocationScale Bijectors" begin @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in [Float64, Float32], @@ -10,12 +25,7 @@ :RepGradELBOStickingTheLanding => RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()), ), - (adbackname, adtype) in Dict( - :ForwarDiff => AutoForwardDiff(), - :ReverseDiff => AutoReverseDiff(), - :Zygote => AutoZygote(), - :Enzyme => AutoEnzyme(), - ) + (adbackname, adtype) in AD_locationscale_bijectors seed = (0x38bef07cf9cc549d) rng = StableRNG(seed) diff --git a/test/runtests.jl b/test/runtests.jl index 2f1d775b..a827dbf9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -47,7 +47,7 @@ end const PROGRESS = haskey(ENV, "PROGRESS") if GROUP == "All" || GROUP == "Inference" - #include("inference/repgradelbo_distributionsad.jl") + include("inference/repgradelbo_distributionsad.jl") include("inference/repgradelbo_locationscale.jl") include("inference/repgradelbo_locationscale_bijectors.jl") end From 5b46a8053f1d1d0034304613a03471dc815086cd Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Wed, 21 Aug 2024 01:07:34 -0400 Subject: [PATCH 14/15] add tests on Julia 1.10 --- .github/workflows/CI.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index b19874b3..224f81d4 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -20,6 +20,7 @@ jobs: matrix: version: - '1.7' + - '1.10' os: - ubuntu-latest - macOS-latest From 18ce79a8a69d8b32c29748a3314ac8aed529ac8d Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Thu, 22 Aug 2024 00:23:19 -0400 Subject: [PATCH 15/15] fix tighten compat bound for Enzyme --- Project.toml | 2 +- test/Project.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index a5bc618b..fff721f2 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ ChainRulesCore = "1.16" DiffResults = "1" Distributions = "0.25.87" DocStringExtensions = "0.8, 0.9" -Enzyme = "0.12.28" +Enzyme = "0.12.32" FillArrays = "1.3" ForwardDiff = "0.10.36" Functors = "0.4" diff --git a/test/Project.toml b/test/Project.toml index ab00716d..251869e7 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -29,7 +29,7 @@ Bijectors = "0.13" DiffResults = "1.0" Distributions = "0.25.100" DistributionsAD = "0.6.45" -Enzyme = "0.12.28" +Enzyme = "0.12.32" FillArrays = "1.6.1" ForwardDiff = "0.10.36" Functors = "0.4.5"