diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index cf5dcdf3f..0f054a6e2 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -16,3 +16,4 @@ steps: setup: julia: - "1.10" + - "1" diff --git a/.github/workflows/Docs.yml b/.github/workflows/Docs.yml index 297438d7a..5b9d7d94e 100644 --- a/.github/workflows/Docs.yml +++ b/.github/workflows/Docs.yml @@ -25,7 +25,7 @@ jobs: - name: Build and deploy Documenter.jl docs uses: TuringLang/actions/DocsDocumenter@main with: - julia-version: 'lts' + julia-version: '1' - name: Run doctests shell: julia --project=docs --color=yes {0} diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml index 732212c62..920bf1c02 100644 --- a/.github/workflows/Format.yml +++ b/.github/workflows/Format.yml @@ -1,8 +1,11 @@ -name: Format suggestions +name: Format on: + push: + branches: + - main pull_request: - + concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. @@ -12,15 +15,7 @@ concurrency: jobs: format: runs-on: ubuntu-latest + steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v2 - with: - version: 1 - - run: | - julia -e 'using Pkg; Pkg.add("JuliaFormatter")' - julia -e 'using JuliaFormatter; format("."; verbose=true)' - - uses: reviewdog/action-suggester@v1 - with: - tool_name: JuliaFormatter - fail_on_error: true + - name: Format code + uses: TuringLang/actions/Format@main diff --git a/Project.toml b/Project.toml index 28d4a3035..5999f0ee1 100644 --- a/Project.toml +++ b/Project.toml @@ -4,10 +4,10 @@ version = "0.7.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -29,23 +29,23 @@ DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" AdvancedVIEnzymeExt = ["Enzyme", "ChainRulesCore"] AdvancedVIMooncakeExt = ["Mooncake", "ChainRulesCore"] AdvancedVIReverseDiffExt = ["ReverseDiff", "ChainRulesCore"] -AdvancedVIDynamicPPLExt = ["DynamicPPL", "Accessors", "Distributions", "DifferentiationInterface", "LogDensityProblems"] +AdvancedVIDynamicPPLExt = ["DynamicPPL", "Accessors", "Distributions", "LogDensityProblems"] [compat] ADTypes = "1" +AbstractPPL = "0.15" Accessors = "0.1" ChainRulesCore = "1" DiffResults = "1" -DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" DocStringExtensions = "0.8, 0.9" -DynamicPPL = "0.40, 0.41" +DynamicPPL = "0.42" Enzyme = "0.13" FillArrays = "1.3" Functors = "0.4, 0.5" LinearAlgebra = "1" LogDensityProblems = "2" -Mooncake = "0.4, 0.5" +Mooncake = "0.5.31" Optimisers = "0.2.16, 0.3, 0.4" ProgressMeter = "1.6" Random = "1" diff --git a/bench/Project.toml b/bench/Project.toml index 4988bba16..a5fc803e1 100644 --- a/bench/Project.toml +++ b/bench/Project.toml @@ -2,6 +2,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" @@ -15,12 +16,12 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] ADTypes = "1" AdvancedVI = "0.7, 0.6" BenchmarkTools = "1" +DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" DistributionsAD = "0.6" Enzyme = "0.13.7" @@ -34,5 +35,4 @@ Random = "1" ReverseDiff = "1" SimpleUnPack = "1" StableRNGs = "1" -Zygote = "0.6, 0.7" julia = "1.10, 1.11.2" diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index a2819bd7d..4dd9d4a30 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -3,7 +3,8 @@ using AdvancedVI using BenchmarkTools using Distributions using DistributionsAD -using Enzyme, ForwardDiff, ReverseDiff, Zygote, Mooncake +using DifferentiationInterface # provides AbstractPPL.prepare for non-Mooncake, non-ForwardDiff backends +using Enzyme, ForwardDiff, Mooncake, ReverseDiff using FillArrays using InteractiveUtils using LinearAlgebra @@ -68,9 +69,8 @@ begin ("RepGradELBO + STL", StickingTheLandingEntropy()), ], (adname, adtype) in [ - ("Zygote", AutoZygote()), + ("Mooncake", AutoMooncake()), ("ReverseDiff", AutoReverseDiff()), - ("Mooncake", AutoMooncake(; config=Mooncake.Config())), # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), ], (familyname, family) in [ diff --git a/docs/Project.toml b/docs/Project.toml index 7a32b4130..ddd05fdd8 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,10 +1,10 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -12,23 +12,22 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" -LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256" OpenML = "8b6db2d4-7670-4922-a472-f9537c81ab66" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" -ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" [compat] ADTypes = "1" +AbstractPPL = "0.15" Accessors = "0.1" AdvancedVI = "0.7, 0.6" -Bijectors = "0.13.6, 0.14, 0.15" +Bijectors = "0.15, 0.16" DataFrames = "1" -DifferentiationInterface = "0.7" Distributions = "0.25" Documenter = "1" FillArrays = "1" @@ -36,13 +35,12 @@ ForwardDiff = "0.10, 1" Functors = "0.5" JSON = "0.21, 1" LogDensityProblems = "2.1.1" -LogDensityProblemsAD = "1" +Mooncake = "0.5.31" NormalizingFlows = "0.2.2" OpenML = "0.3" Optimisers = "0.3, 0.4" Plots = "1" QuasiMonteCarlo = "0.3" -ReverseDiff = "1" StanLogDensityProblems = "0.1" StatsFuns = "1, 2" julia = "1.10, 1.11.2" diff --git a/docs/make.jl b/docs/make.jl index f85540832..44acd186d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,4 +1,9 @@ +# Force GR (Plots.jl's default backend) into its no-display workstation so +# `savefig` writes files but no interactive window is opened. Must be set +# before GR loads. +ENV["GKSwstype"] = "100" + using AdvancedVI using Documenter diff --git a/docs/src/families.md b/docs/src/families.md index a7f2e19b8..57dee3333 100644 --- a/docs/src/families.md +++ b/docs/src/families.md @@ -138,7 +138,7 @@ using LinearAlgebra using LogDensityProblems using Optimisers using Plots -using ForwardDiff, ReverseDiff +using ForwardDiff, Mooncake struct Target{D} dist::D @@ -184,7 +184,7 @@ D = ones(n_dims) U = zeros(n_dims, 3) q0_lr = LowRankGaussian(μ, D, U) -alg = KLMinRepGradDescent(AutoReverseDiff(); optimizer=Adam(0.01), operator=ClipScale()) +alg = KLMinRepGradDescent(AutoMooncake(); optimizer=Adam(0.01), operator=ClipScale()) max_iter = 10^4 diff --git a/docs/src/klminrepgraddescent.md b/docs/src/klminrepgraddescent.md index 771cf99cd..1c75e55f1 100644 --- a/docs/src/klminrepgraddescent.md +++ b/docs/src/klminrepgraddescent.md @@ -103,7 +103,7 @@ using Plots using Random using Optimisers -using ADTypes, ForwardDiff, ReverseDiff +using ADTypes, ForwardDiff, Mooncake using AdvancedVI struct Dist{D} @@ -150,10 +150,7 @@ Recall that the original ADVI objective with a closed-form entropy (CFE) is give n_montecarlo = 16; cfe = KLMinRepGradDescent( - AutoReverseDiff(); - entropy=ClosedFormEntropy(), - optimizer=Adam(1e-2), - operator=ClipScale(), + AutoMooncake(); entropy=ClosedFormEntropy(), optimizer=Adam(1e-2), operator=ClipScale() ) nothing ``` @@ -162,7 +159,7 @@ The repgradelbo estimator can instead be created as follows: ```@example repgradelbo stl = KLMinRepGradDescent( - AutoReverseDiff(); + AutoMooncake(); entropy=StickingTheLandingEntropy(), optimizer=Adam(1e-2), operator=ClipScale(), @@ -262,6 +259,10 @@ In this case, it suffices to override its `rand` specialization as follows: using QuasiMonteCarlo using StatsFuns +# QMC samples are inputs to AD, not parameters; declare them non-differentiable +# so Mooncake doesn't trace through Sobol/Owen bit-twiddling. +Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(QuasiMonteCarlo.sample),Vararg} + qmcrng = SobolSample(; R=OwenScramble(; base=2, pad=32)) function Distributions.rand( @@ -281,7 +282,7 @@ nothing ```@setup repgradelbo _, info_qmc, _ = AdvancedVI.optimize( - KLMinRepGradDescent(AutoReverseDiff(); n_samples=n_montecarlo, optimizer=Adam(1e-2), operator=ClipScale()), + KLMinRepGradDescent(AutoMooncake(); n_samples=n_montecarlo, optimizer=Adam(1e-2), operator=ClipScale()), max_iter, model, q0; diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index 31e57333e..5fb315918 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -143,10 +143,10 @@ nothing For the VI algorithm, we will use `KLMinRepGradDescent`: ```@example basic -using ADTypes, ReverseDiff +using ADTypes, Mooncake using AdvancedVI -alg = KLMinRepGradDescent(ADTypes.AutoReverseDiff(); operator=ClipScale()); +alg = KLMinRepGradDescent(ADTypes.AutoMooncake(); operator=ClipScale()); nothing ``` @@ -159,19 +159,7 @@ Location-scale family distributions require the scale matrix to have strictly po Here, the projection operator `ClipScale` ensures this. `KLMinRepGradDescent`, in particular, assumes that the target `LogDensityProblem` is differentiable. -If the `LogDensityProblem` has a differentiation [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/dev/#LogDensityProblems.capabilities) of at least first-order, we can take advantage of this. -For this example, we will use `LogDensityProblemsAD` to equip our problem with a first-order capability: - -[^TL2014]: Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In *International Conference on Machine Learning*. PMLR. -[^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. -[^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. -```@example basic -using DifferentiationInterface: DifferentiationInterface -using LogDensityProblemsAD: LogDensityProblemsAD - -prob_trans_ad = LogDensityProblemsAD.ADgradient(ADTypes.AutoForwardDiff(), prob_trans) -nothing -``` +If the problem provides only a zeroth-order [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/dev/#LogDensityProblems.capabilities), `AdvancedVI` will differentiate through `LogDensityProblems.logdensity` directly using the AD backend supplied to the algorithm. For the variational family, we will consider a `FullRankGaussian` approximation: @@ -190,7 +178,7 @@ We can now run VI: ```@example basic max_iter = 10^4 -q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans_ad, q; show_progress=false) +q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans, q; show_progress=false) nothing ``` @@ -286,7 +274,7 @@ The `callback` can be supplied to `optimize`: ```@example basic max_iter = 10^4 q_out, info, _ = AdvancedVI.optimize( - alg, max_iter, prob_trans_ad, q; show_progress=false, callback=callback + alg, max_iter, prob_trans, q; show_progress=false, callback=callback ) nothing ``` @@ -331,3 +319,7 @@ nothing ![](basic_example_acc.svg) Clearly, the accuracy is improving over time. + +[^TL2014]: Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In *International Conference on Machine Learning*. PMLR. +[^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. +[^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. diff --git a/docs/src/tutorials/constrained.md b/docs/src/tutorials/constrained.md index c3b3fb781..ca7462c59 100644 --- a/docs/src/tutorials/constrained.md +++ b/docs/src/tutorials/constrained.md @@ -147,21 +147,41 @@ This approach only requires the user to implement the model-specific `Bijectors. The rest can be done by simply copy-pasting the code below: ```@example constraints -struct TransformedLogDensityProblem{Prob,BInv} +using ADTypes +using AbstractPPL: AbstractPPL +using Mooncake: Mooncake + +struct TransformedLogDensityProblem{Prob,BInv,Prep} prob::Prob binv::BInv + prep::Prep end -function TransformedLogDensityProblem(prob) +function TransformedLogDensityProblem(prob; adtype::ADTypes.AbstractADType=AutoMooncake()) b = Bijectors.bijector(prob) binv = Bijectors.inverse(b) - return TransformedLogDensityProblem{typeof(prob),typeof(binv)}(prob, binv) + d_unc = prod(Bijectors.output_size(b, (LogDensityProblems.dimension(prob),))) + + f = let prob = prob, binv = binv + function (θ_trans) + θ, logabsdetjac = Bijectors.with_logabsdet_jacobian(binv, θ_trans) + return LogDensityProblems.logdensity(prob, θ) + logabsdetjac + end + end + prep = AbstractPPL.prepare(adtype, f, zeros(d_unc)) + return TransformedLogDensityProblem(prob, binv, prep) end function LogDensityProblems.logdensity(prob_trans::TransformedLogDensityProblem, θ_trans) - (; prob, binv) = prob_trans - θ, logabsdetjac = Bijectors.with_logabsdet_jacobian(binv, θ_trans) - return LogDensityProblems.logdensity(prob, θ) + logabsdetjac + return prob_trans.prep(θ_trans) +end + +function LogDensityProblems.logdensity_and_gradient( + prob_trans::TransformedLogDensityProblem, θ_trans +) + # `prep` was built for a dense `Vector`; collect view-typed inputs + # (e.g. minibatch column views) to match the AD backend's expectations. + return AbstractPPL.value_and_gradient!!(prob_trans.prep, collect(θ_trans)) end function LogDensityProblems.dimension(prob_trans::TransformedLogDensityProblem) @@ -171,15 +191,14 @@ function LogDensityProblems.dimension(prob_trans::TransformedLogDensityProblem) return prod(Bijectors.output_size(b, (d,))) end -function LogDensityProblems.capabilities( - ::Type{TransformedLogDensityProblem{Prob,BInv}} -) where {Prob,BInv} - return LogDensityProblems.capabilities(Prob) +function LogDensityProblems.capabilities(::Type{<:TransformedLogDensityProblem}) + return LogDensityProblems.LogDensityOrder{1}() end nothing ``` Wrapping `prob` with `TransformedLogDensityProblem` yields our unconstrained posterior. +The `AbstractPPL.prepare` call traces the AD backend once at construction time, so the per-iteration `logdensity_and_gradient` call just runs the cached evaluator. ```@example constraints prob_trans = TransformedLogDensityProblem(prob) @@ -188,17 +207,6 @@ x = randn(LogDensityProblems.dimension(prob_trans)) # sample on an unconstrained LogDensityProblems.logdensity(prob_trans, x) ``` -We can also wrap `prob_trans` with `LogDensityProblemsAD.ADGradient` to make it differentiable. - -```@example constraints -using LogDensityProblemsAD -using ADTypes, ReverseDiff - -prob_trans_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoReverseDiff(; compile=true), prob_trans; x=randn(2) -) -``` - Let's now run VI to verify that it works. Here, we will use `FisherMinBatchMatch`, which expects an unconstrained posterior. @@ -206,11 +214,11 @@ Here, we will use `FisherMinBatchMatch`, which expects an unconstrained posterio using AdvancedVI using LinearAlgebra -d = LogDensityProblems.dimension(prob_trans_ad) +d = LogDensityProblems.dimension(prob_trans) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) q_opt, info, _ = AdvancedVI.optimize( - FisherMinBatchMatch(), 100, prob_trans_ad, q; show_progress=false + FisherMinBatchMatch(), 100, prob_trans, q; show_progress=false ) nothing ``` diff --git a/docs/src/tutorials/flows.md b/docs/src/tutorials/flows.md index 4a408160b..4da07546c 100644 --- a/docs/src/tutorials/flows.md +++ b/docs/src/tutorials/flows.md @@ -138,14 +138,9 @@ Let's instantiate the model: ```@example flow using ADTypes: ADTypes -using ReverseDiff: ReverseDiff -using DifferentiationInterface: DifferentiationInterface -using LogDensityProblemsAD: LogDensityProblemsAD +using Mooncake: Mooncake prob_trans = TransformedLogDensityProblem(prob) -prob_trans_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoForwardDiff(), prob_trans; x=[1.0, 1.0] -) nothing ``` @@ -155,12 +150,12 @@ For the algorithm, we will use the `KLMinRepGradProxDescent` objective. using AdvancedVI using LinearAlgebra -d = LogDensityProblems.dimension(prob_trans_ad) +d = LogDensityProblems.dimension(prob_trans) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(I, d, d))) max_iter = 3 * 10^3 -alg = KLMinRepGradProxDescent(ADTypes.AutoReverseDiff(; compile=true)) -q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans_ad, q; show_progress=false) +alg = KLMinRepGradProxDescent(ADTypes.AutoMooncake()) +q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans, q; show_progress=false) b = Bijectors.bijector(prob) binv = Bijectors.inverse(b) q_out_trans = Bijectors.TransformedDistribution(q_out, binv) @@ -229,7 +224,7 @@ Furthermore, Agrawal *et al.*[^AD2025] claim that using a larger number of Monte using Optimisers: Optimisers alg_flow = KLMinRepGradDescent( - ADTypes.AutoReverseDiff(; compile=true); + ADTypes.AutoMooncake(); n_samples=8, optimizer=Optimisers.Adam(1e-2), operator=IdentityOperator(), @@ -243,7 +238,7 @@ Without further due, let's now run VI: ```@example flow max_iter = 300 q_flow_out, info_flow, _ = AdvancedVI.optimize( - alg_flow, max_iter, prob_trans_ad, q_flow; show_progress=false + alg_flow, max_iter, prob_trans, q_flow; show_progress=false ) q_flow_out_trans = Bijectors.TransformedDistribution(q_flow_out, binv) nothing diff --git a/docs/src/tutorials/stan.md b/docs/src/tutorials/stan.md index 966a65908..50edb3865 100644 --- a/docs/src/tutorials/stan.md +++ b/docs/src/tutorials/stan.md @@ -75,13 +75,13 @@ nothing The rest is the same as all `LogDensityProblem` with the exception of how to deal with constrainted variables: Since `StanLogDensityProblems` automatically transforms the support of the target problem to be unconstrained, we do not need to involve `Bijectors`. ```@example stan -using ADTypes, ReverseDiff +using ADTypes, Mooncake using AdvancedVI using LinearAlgebra using LogDensityProblems using Plots -alg = KLMinRepGradDescent(ADTypes.AutoReverseDiff(); operator=ClipScale()) +alg = KLMinRepGradDescent(ADTypes.AutoMooncake(); operator=ClipScale()) d = LogDensityProblems.dimension(model) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.37 * I, d, d))) diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index 8d33239fb..294781e63 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -80,34 +80,25 @@ X = hcat(X, ones(size(X, 1))) nothing ``` -Let's now istantiate the model and set up automatic differentiation using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl?tab=readme-ov-file). +Let's now instantiate the model. ```@example subsampling -using ADTypes, ReverseDiff -using LogDensityProblemsAD +using ADTypes, Mooncake prob = LogReg(X, y, size(X, 1)) -prob_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoReverseDiff(), prob; x=randn(LogDensityProblems.dimension(prob)) -) nothing ``` To enable subsampling, `LogReg` has to implement the method `AdvancedVI.subsample`. For our model, this is fairly simple: We only need to select the rows of `X` and the elements of `y` corresponding to the batch of data points. -As subtle point here is that we wrapped `prob` with `LogDensityProblemsAD.ADgradient` into `prob_ad`. -Therefore, `AdvancedVI` sees `prob_ad` and not `prob`. -This means we have to specialize `AdvancedVI.subsample` to `typeof(prob_ad)` and not `LogReg`. ```@example subsampling using Accessors using AdvancedVI -function AdvancedVI.subsample(prob::typeof(prob_ad), idx) - (; X, y, n_data) = parent(prob_ad) - prob′ = @set prob.ℓ.X = X[idx, :] - prob′′ = @set prob′.ℓ.y = y[idx] - return prob′′ +function AdvancedVI.subsample(prob::LogReg, idx) + prob′ = @set prob.X = prob.X[idx, :] + return @set prob′.y = prob′.y[idx] end nothing ``` @@ -129,7 +120,7 @@ We will us a batch size of 32, which results in `313 = length(subsampling) = cei dataset = 1:size(prob.X, 1) batchsize = 32 subsampling = ReshufflingBatchSubsampling(dataset, batchsize) -alg_sub = KLMinRepGradProxDescent(ADTypes.AutoReverseDiff(; compile=true); subsampling) +alg_sub = KLMinRepGradProxDescent(ADTypes.AutoMooncake(); subsampling) nothing ``` @@ -148,7 +139,7 @@ nothing If we don't supply a subsampling strategy to `KLMinRepGradProxDescent`, subsampling will not be used. ```@example subsampling -alg_full = KLMinRepGradProxDescent(ADTypes.AutoReverseDiff(; compile=true)) +alg_full = KLMinRepGradProxDescent(ADTypes.AutoMooncake()) nothing ``` @@ -157,7 +148,7 @@ The variational family will be set up as follows: ```@example subsampling using LinearAlgebra -d = LogDensityProblems.dimension(prob_ad) +d = LogDensityProblems.dimension(prob) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) nothing ``` @@ -208,12 +199,12 @@ end time_begin = time() _, info_full, _ = AdvancedVI.optimize( - alg_full, max_iter, prob_ad, q; show_progress=false, callback + alg_full, max_iter, prob, q; show_progress=false, callback ); time_begin = time() _, info_sub, _ = AdvancedVI.optimize( - alg_sub, max_iter, prob_ad, q; show_progress=false, callback + alg_sub, max_iter, prob, q; show_progress=false, callback ); nothing ``` diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index eaff70ae7..687a36bb0 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -1,10 +1,8 @@ module AdvancedVIDynamicPPLExt using ADTypes: ADTypes -using Accessors using AdvancedVI: AdvancedVI -using DifferentiationInterface: DifferentiationInterface -using Distributions: Distributions +using AbstractPPL: AbstractPPL using DynamicPPL: DynamicPPL using LogDensityProblems: LogDensityProblems using Random @@ -15,47 +13,25 @@ function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) return LogDensityProblems.LogDensityOrder{1}() end -function adtype_capabilities( - ::Type{ - <:Union{ - <:ADTypes.AutoForwardDiff, - <:ADTypes.AutoReverseDiff, - <:ADTypes.AutoMooncake, - <:ADTypes.AutoEnzyme, - <:DifferentiationInterface.SecondOrder, - }, - }, -) - return LogDensityProblems.LogDensityOrder{2}() -end - -struct DynamicPPLModelLogDensityFunction{ - Model<:DynamicPPL.Model, - LogLikeAdj<:Real, - VarInfo<:DynamicPPL.AbstractVarInfo, - ADType<:ADTypes.AbstractADType, - PrepGrad<:Union{Nothing,DifferentiationInterface.GradientPrep}, - PrepHess<:Union{Nothing,DifferentiationInterface.HessianPrep}, -} - model::Model - loglikeadj::LogLikeAdj - varinfo::VarInfo - adtype::ADType - prep_grad::PrepGrad - prep_hess::PrepHess +# `getlogdensity` callable for `DynamicPPL.logdensity_internal`: reads the +# current `loglikeadj` through a Ref so the mutation done by `subsample` is +# observed without rebuilding any AD prep. +struct WeightedLogJoint{R<:Base.RefValue{<:Real}} + loglikeadj_ref::R end - -function logdensity_impl( - params, model::DynamicPPL.Model, loglikeadj::Real, varinfo::DynamicPPL.AbstractVarInfo -) - vi = DynamicPPL.unflatten!!(varinfo, params) - _, vi = DynamicPPL.evaluate!!(model, vi) +function (g::WeightedLogJoint)(vi) loglike = DynamicPPL.getloglikelihood(vi) logprior = DynamicPPL.getlogprior(vi) logjac = DynamicPPL.getlogjac(vi) - return convert(eltype(params), loglikeadj) * loglike + logprior - logjac + return g.loglikeadj_ref[] * loglike + logprior - logjac end +const _DEFAULT_LDF_ACCS = DynamicPPL.AccumulatorTuple(( + DynamicPPL.LogPriorAccumulator(), + DynamicPPL.LogJacobianAccumulator(), + DynamicPPL.LogLikelihoodAccumulator(), +)) + function subsample_dynamicpplmodel( model::DynamicPPL.Model{F,A,D,M,Ta,Td,Ctx,Threaded}, batch ) where {F,A,D,M,Ta,Td,Ctx,Threaded} @@ -63,6 +39,36 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end +# `model` is the original (unsubsampled) source of truth; `subsample` must read +# it (not `model_ref[]`) to get the full-dataset length on every call. +# `model_ref`/`loglikeadj_ref` are mutated in place by `subsample` so the +# closure inside `prep_grad`/`prep_hess` stays valid across subsampling steps. +# `model_ref` is `Ref{Any}` because `subsample_dynamicpplmodel`'s output type +# varies with the batch (a typed Ref would throw on reassignment), and because +# compiled-tape backends would otherwise bake the deref into the tape and miss +# the `subsample` update. +struct DynamicPPLModelLogDensityFunction{ + Model<:DynamicPPL.Model, + LogLikeAdj<:Real, + Ranges<:DynamicPPL.VarNamedTuple, + Strategy<:DynamicPPL.AbstractTransformStrategy, + GetLogDensity, + ADType<:Union{Nothing,ADTypes.AbstractADType}, + PrepGrad, + PrepHess, +} + model::Model + model_ref::Ref{Any} + loglikeadj_ref::Ref{LogLikeAdj} + ranges_and_transforms::Ranges + transform_strategy::Strategy + getlogdensity::GetLogDensity + adtype::ADType + prep_grad::PrepGrad + prep_hess::PrepHess + dim::Int +end + function DynamicPPLModelLogDensityFunction( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; @@ -80,32 +86,36 @@ function DynamicPPLModelLogDensityFunction( subsample_dynamicpplmodel(model, batch) end - params = [val for val in varinfo[:]] - cap = adtype_capabilities(typeof(adtype)) - prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() - DifferentiationInterface.prepare_gradient( - logdensity_impl, - DifferentiationInterface.inner(adtype), + ranges_and_transforms, params = DynamicPPL.get_rat_and_samplevec(varinfo.values) + transform_strategy = DynamicPPL.infer_transform_strategy_from_values( + ranges_and_transforms + ) + + model_ref = Ref{Any}(model_sub) + loglikeadj_ref = Ref(float(loglikeadj)) + getlogdensity = WeightedLogJoint(loglikeadj_ref) + f = + params -> DynamicPPL.logdensity_internal( params, - DifferentiationInterface.Constant(model_sub), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), + model_ref[], + getlogdensity, + ranges_and_transforms, + transform_strategy, + _DEFAULT_LDF_ACCS, ) + cap = adtype_capabilities(typeof(adtype)) + + prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() + AbstractPPL.prepare(adtype, f, params) else nothing end - prep_hess = if cap >= LogDensityProblems.LogDensityOrder{2}() && use_hessian + prep_hess = if cap >= LogDensityProblems.LogDensityOrder{1}() && use_hessian try - DifferentiationInterface.prepare_hessian( - logdensity_impl, - adtype, - params, - DifferentiationInterface.Constant(model_sub), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), - ) - catch - @warn "The selected AD backend has second-order capabilities but `DifferentiationInterface.prepare_hessian` failed. AdvancedVI will treat the model to only have first-order capability." + AbstractPPL.prepare(adtype, f, params; order=2) + catch err + err isa MethodError || rethrow() + @warn "The selected AD backend does not support `AbstractPPL.prepare(...; order=2)`. AdvancedVI will treat the model as first-order only." nothing end else @@ -113,74 +123,75 @@ function DynamicPPLModelLogDensityFunction( end return DynamicPPLModelLogDensityFunction{ typeof(model), - typeof(loglikeadj), - typeof(varinfo), + eltype(loglikeadj_ref), + typeof(ranges_and_transforms), + typeof(transform_strategy), + typeof(getlogdensity), typeof(adtype), typeof(prep_grad), typeof(prep_hess), }( - model, loglikeadj, varinfo, adtype, prep_grad, prep_hess + model, + model_ref, + loglikeadj_ref, + ranges_and_transforms, + transform_strategy, + getlogdensity, + adtype, + prep_grad, + prep_hess, + length(params), ) end function LogDensityProblems.logdensity(prob::DynamicPPLModelLogDensityFunction, params) - (; model, loglikeadj, varinfo) = prob - return logdensity_impl(params, model, loglikeadj, varinfo) + return DynamicPPL.logdensity_internal( + params, + prob.model_ref[], + prob.getlogdensity, + prob.ranges_and_transforms, + prob.transform_strategy, + _DEFAULT_LDF_ACCS, + ) end +# `!!` may alias internal buffers of `prep_*`; copy so callers can retain the +# arrays past the next AD call. function LogDensityProblems.logdensity_and_gradient( prob::DynamicPPLModelLogDensityFunction, params ) - (; model, adtype, loglikeadj, varinfo, prep_grad) = prob - return DifferentiationInterface.value_and_gradient( - logdensity_impl, - prep_grad, - DifferentiationInterface.inner(adtype), - params, - DifferentiationInterface.Constant(model), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), - ) + val, grad = AbstractPPL.value_and_gradient!!(prob.prep_grad, params) + return val, copy(grad) end function LogDensityProblems.logdensity_gradient_and_hessian( prob::DynamicPPLModelLogDensityFunction, params ) - (; model, adtype, loglikeadj, varinfo, prep_hess) = prob - return DifferentiationInterface.value_gradient_and_hessian( - logdensity_impl, - prep_hess, - adtype, - params, - DifferentiationInterface.Constant(model), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), - ) + val, grad, H = AbstractPPL.value_gradient_and_hessian!!(prob.prep_hess, params) + return val, copy(grad), copy(H) end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,L,V,ADType,PG,PH}} -) where {M,L,V,ADType<:ADTypes.AbstractADType,PG,PH} - return if PH != Nothing + ::Type{<:DynamicPPLModelLogDensityFunction{Model,L,R,S,G,ADType,PG,PH}} +) where {Model,L,R,S,G,ADType<:Union{Nothing,ADTypes.AbstractADType},PG,PH} + return if PH !== Nothing LogDensityProblems.LogDensityOrder{2}() - elseif PG != Nothing + elseif PG !== Nothing LogDensityProblems.LogDensityOrder{1}() else LogDensityProblems.LogDensityOrder{0}() end end -function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) - return length(prob.varinfo[:]) -end +LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) = prob.dim function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) - model = prob.model + model = prob.model # full dataset — `model_ref[]` would already be subsampled if !haskey(model.defaults, :datapoints) throw( ArgumentError( - "Subsampling is turned on, but the model does not have have a `datapoints` keyword argument.", + "Subsampling is turned on, but the model does not have a `datapoints` keyword argument.", ), ) end @@ -188,11 +199,13 @@ function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) n_datapoints = length(model.defaults.datapoints) batchsize = length(batch) model_sub = subsample_dynamicpplmodel(model, batch) - loglikeadj = n_datapoints / batchsize + T = eltype(prob.loglikeadj_ref) + loglikeadj = T(n_datapoints) / T(batchsize) + + prob.model_ref[] = model_sub + prob.loglikeadj_ref[] = loglikeadj - prob′ = @set prob.model = model_sub - prob′′ = @set prob′.loglikeadj = loglikeadj - return prob′′ + return prob end end diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 25658c935..727fb2f5d 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -17,7 +17,7 @@ using LogDensityProblems using ADTypes using DiffResults -using DifferentiationInterface +using AbstractPPL: AbstractPPL using ChainRulesCore: ChainRulesCore using FillArrays @@ -33,43 +33,41 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif `f` may receive auxiliary input as `f(x,aux)`. # Arguments -- `ad::ADTypes.AbstractADType`: - automatic differentiation backend. Currently supports - `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, - `ADTypes.AutoMooncake()` and - `ADTypes.AutoEnzyme(; - mode=Enzyme.set_runtime_activity(Enzyme.Reverse), - function_annotation=Enzyme.Const, - )`. - If one wants to use `AutoEnzyme`, please make sure to include the `set_runtime_activity` and `function_annotation` as shown above. +- `ad::ADTypes.AbstractADType`: automatic differentiation backend; any backend + with an `AbstractPPL.prepare` method is supported. `AutoEnzyme` requires + `set_runtime_activity` and `function_annotation=Enzyme.Const`. + `AutoReverseDiff(; compile=true)` is rejected because compiled tapes + freeze captured values at prep time. - `f`: Function subject to differentiation. - `x`: The point to evaluate the gradient. - `aux`: Auxiliary input passed to `f`. -- `prep`: Output of `DifferentiationInterface.prepare_gradient`. +- `prep`: Output of `_prepare_gradient`. - `out::DiffResults.MutableDiffResult`: Buffer to contain the output gradient and function value. """ function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, ad::ADTypes.AbstractADType, x, aux ) - grad_buf = DiffResults.gradient(out) - y, _ = DifferentiationInterface.value_and_gradient!(f, grad_buf, ad, x, Constant(aux)) - DiffResults.value!(out, y) + prepared = AbstractPPL.prepare(ad, Base.Fix2(f, aux), x) + val, grad = AbstractPPL.value_and_gradient!!(prepared, x) + DiffResults.value!(out, val) + copyto!(DiffResults.gradient(out), grad) return out end function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, prep, ad::ADTypes.AbstractADType, x, aux ) - grad_buf = DiffResults.gradient(out) - y, _ = DifferentiationInterface.value_and_gradient!( - f, grad_buf, prep, ad, x, Constant(aux) - ) - DiffResults.value!(out, y) + # `context[1]` is the `Ref(aux)` placed there by `_prepare_gradient`; mutating + # it lets a cached prep see the new `aux` without rebuilding AD machinery. + prep.evaluator.context[1][] = aux + val, grad = AbstractPPL.value_and_gradient!!(prep, x) + DiffResults.value!(out, val) + copyto!(DiffResults.gradient(out), grad) return out end """ - _prepare_gradient!(f, ad, x, aux) + _prepare_gradient(f, ad, x, aux) Prepare AD backend for taking gradients of a function `f` at `x` using the automatic differentiation backend `ad`. @@ -80,7 +78,23 @@ Prepare AD backend for taking gradients of a function `f` at `x` using the autom - `aux`: Auxiliary input passed to `f`. """ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) - return DifferentiationInterface.prepare_gradient(f, ad, x, Constant(aux)) + return AbstractPPL.prepare(ad, (x, aref) -> f(x, aref[]), x; context=(Ref(aux),)) +end + +# Compiled-tape ReverseDiff bakes context values into the tape at prep time, +# so the `Ref(aux)` trick above would feed stale `aux` to AD. Reject it +# loudly rather than silently producing wrong gradients. +function _prepare_gradient(::Any, ::ADTypes.AutoReverseDiff{true}, ::Any, ::Any) + throw( + ArgumentError( + "`AutoReverseDiff(; compile=true)` is not safe to use with AdvancedVI: " * + "compiled tapes freeze captured values at preparation time, so " * + "subsequent iterations differentiate against stale data and silently " * + "produce incorrect gradients. Use `AutoReverseDiff(; compile=false)` " * + "(or another reverse-mode backend such as `AutoMooncake` or " * + "`AutoEnzyme`) instead.", + ), + ) end """ @@ -292,6 +306,9 @@ export estimate_objective Inform `model` or `q` to only use the data points designated by the iterable collection `batch`. For `model`, the log-density should also be adjusted to account for the change in number of data points. + +Implementations may mutate `model`/`q` in place and return the same object, +so callers must not assume the returned value is distinct from the input. """ subsample(model_or_q::Any, ::Any) = model_or_q diff --git a/src/algorithms/subsampledobjective.jl b/src/algorithms/subsampledobjective.jl index 21734f7ca..ad1f858ef 100644 --- a/src/algorithms/subsampledobjective.jl +++ b/src/algorithms/subsampledobjective.jl @@ -32,7 +32,7 @@ function init( sub_st = init(rng, subsampling) # This is necessary to ensure that `init` sees the type "conditioned" on a minibatch - # when calling `DifferentiationInterface.prepare_*` inside it. + # so that any prepared AD evaluator inside it sees the correct batch-subsampled type. batch, _, _ = step(rng, subsampling, sub_st, true) prob_sub = subsample(prob, batch) q_init_sub = subsample(q_init, batch) diff --git a/src/families/location_scale_low_rank.jl b/src/families/location_scale_low_rank.jl index f3e444a30..3414a5e45 100644 --- a/src/families/location_scale_low_rank.jl +++ b/src/families/location_scale_low_rank.jl @@ -36,7 +36,8 @@ function StatsBase.entropy(q::MvLocationScaleLowRank) (; location, scale_diag, scale_factors, dist) = q n_dims = length(location) scale_diag2 = scale_diag .* scale_diag - UtDinvU = Hermitian(scale_factors' * (scale_factors ./ scale_diag2)) + # `Symmetric` rather than `Hermitian`: Mooncake lacks an `rrule` for the `Hermitian` path. + UtDinvU = Symmetric(scale_factors' * (scale_factors ./ scale_diag2)) logdetΣ = 2 * sum(log.(scale_diag)) + logdet(I + UtDinvU) return n_dims * convert(eltype(location), entropy(dist)) + logdetΣ / 2 end diff --git a/src/reshuffling.jl b/src/reshuffling.jl index e0e50cfbd..dd7f95aa8 100644 --- a/src/reshuffling.jl +++ b/src/reshuffling.jl @@ -49,7 +49,7 @@ function step( # Ignore the trailing batch if its size is smaller than `batchsize`. # This should only be used when estimating gradients during optimization. # This is necessary to ensure that all batches have the same size. - # Otherwise, `DifferentiationInterface.prepare_*` behaves incorrectly. + # Otherwise, prepared AD evaluators may see inconsistent batch sizes. (sub_step, batch), iterator = Iterators.peel(iterator) end epoch = epoch + 1 diff --git a/test/Project.toml b/test/Project.toml index 356cbf4e1..7cb2e8a7a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -27,7 +27,7 @@ ADTypes = "0.2.1, 1" DiffResults = "1" DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" -DynamicPPL = "0.40, 0.41" +DynamicPPL = "0.42" Enzyme = "0.13, 0.14, 0.15" FillArrays = "1.6.1" ForwardDiff = "0.10.36, 1" diff --git a/test/runtests.jl b/test/runtests.jl index b50e69bfc..464d7e1c9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,6 +17,7 @@ using Random, StableRNGs using Statistics using StatsBase +using DifferentiationInterface # triggers DI's Enzyme extension using AdvancedVI const PROGRESS = haskey(ENV, "PROGRESS") @@ -29,7 +30,7 @@ elseif AD_str == "ForwardDiff" AutoForwardDiff() elseif AD_str == "Mooncake" using Mooncake - AutoMooncake(; config=Mooncake.Config()) + AutoMooncake() elseif AD_str == "Enzyme" using Enzyme AutoEnzyme(;