Update to the AdvancedVI@0.3 interface #2506

Red-Portal · 2025-03-14T23:14:25Z

This PR aims to update Turing's Variational module to match AdvancedVI's new interface starting from v0.3. I will try not to change the interface too much, but given the new features in AdvancedVI, I think breaking changes will be inevitable. Though the focus will be to provide a good default setting rather than to expose all the features.

Currently proposed interface:

using Turing
using AdvancedVI

d = randn(100)
Turing.@model function model()
   x ~ MvNormal(randn(100), 1)
   y ~ InverseGamma()
end
m = model()
q = Turing.Variational.fullrank_gaussian(Random.default_rng(), m)
n_iters = 1000
Turing.Variational.vi(m, q, n_iters)

Closes #2507
Closes #2508
Closes #2430

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

…nto update_advancedvi

src/variational/VariationalInference.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

…nto update_advancedvi

yebai · 2025-03-20T21:21:36Z

@Red-Portal, can you fix the tests before I take a look?

Red-Portal · 2025-03-20T21:31:09Z

@yebai I marked the PR as a draft so that we can first agree on an interface, and then I flesh out the implementation and the tests. Do we wish we proceed in another way?

yebai · 2025-03-20T21:51:41Z

Let's address the interface later or in a separate PR since that might require more discussions. For this PR, let's try to keep the VI interface non-breaking where possible.

…vancedvi

…nto update_advancedvi

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

yebai · 2025-04-18T20:18:33Z

@Red-Portal can you take a look at the following error:

ERROR: LoadError: UndefVarError: `turnprogress` not defined in `AdvancedVI`

codecov · 2025-04-21T11:03:58Z

Codecov Report

Attention: Patch coverage is 0% with 89 lines in your changes missing coverage. Please review.

Project coverage is 55.13%. Comparing base (afb5c44) to head (231d6e2).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/variational/VariationalInference.jl	0.00%	57 Missing ⚠️
src/variational/bijectors.jl	0.00%	32 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (afb5c44) and HEAD (231d6e2). Click for more details.

HEAD has 12 uploads less than BASE

Flag BASE (afb5c44) HEAD (231d6e2)

42 30

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2506       +/-   ##
===========================================
- Coverage   86.04%   55.13%   -30.92%     
===========================================
  Files          21       21               
  Lines        1455     1471       +16     
===========================================
- Hits         1252      811      -441     
- Misses        203      660      +457

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

yebai

Thanks @Red-Portal. I left some comments below.

One high-level comment: I suggest we unify vi_fullrank_gaussian and vi_meanfield_gaussian into a single function, q_distribution(...; gaussian=true|false, fullrank=true|false), to reduce code redundancy.

Note on CI errors:

CI compains about missing ADVI() and TruncatedADAGrad().

yebai · 2025-04-21T11:01:30Z

src/variational/VariationalInference.jl

+
+        if isfinite(energy)
+            return scale
+        elseif n_trial == num_max_trials


Suggested change

elseif n_trial == num_max_trials

else

yebai · 2025-04-21T11:03:18Z

src/variational/VariationalInference.jl

+    end
+end
+
+function meanfield_gaussian(


Suggested change

function meanfield_gaussian(

function q_meanfield_gaussian(

yebai · 2025-04-21T11:03:34Z

src/variational/VariationalInference.jl

+    return Bijectors.transformed(q, Bijectors.inverse(b))
+end
+
+function meanfield_gaussian(


Suggested change

function meanfield_gaussian(

function q_meanfield_gaussian(

yebai · 2025-04-21T11:03:50Z

src/variational/VariationalInference.jl

+    return meanfield_gaussian(Random.default_rng(), model; location, scale, kwargs...)
+end
+
+function fullrank_gaussian(


Suggested change

function fullrank_gaussian(

function q_fullrank_gaussian(

yebai · 2025-04-21T11:05:36Z

src/variational/VariationalInference.jl

+    # Use linked `varinfo` to determine the correct number of parameters.
+    # TODO: Replace with `length` once this is implemented for `VarInfo`.
+    varinfo_linked = DynamicPPL.link(varinfo, model)
+    num_params = length(varinfo_linked[:])


Can we get the dimentionality via num_params = length(varinfo_linked) instead of length(varinfo_linked[:])?

cc @mkarikom

yebai · 2025-04-21T11:06:31Z

src/variational/VariationalInference.jl

 end

-# VI algorithms
-include("advi.jl")
+function fullrank_gaussian(


Suggested change

function fullrank_gaussian(

function q_fullrank_gaussian(

yebai · 2025-04-21T11:09:46Z

src/variational/bijectors.jl

+    return reshape_outer ∘ f ∘ reshape_inner
+end
+
+"""


Please move Bijectors.bijector(model::DynamicPPL.Model,...) to DynamicPPL.

cc @mhauru

yebai · 2025-04-21T11:11:43Z

src/variational/bijectors.jl

+Wraps a bijector `f` such that it operates on vectors of length `prod(in_size)` and produces
+a vector of length `prod(Bijectors.output(f, in_size))`.
+"""
+function wrap_in_vec_reshape(f, in_size)


Since this function is only used once, I suggest we inline it and add comments to explain its behaviour.

Red-Portal and others added 12 commits March 14, 2025 19:09

update to match the AdvancedVI@0.3 interface

ed6946c

run formatter

a94269d

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

a4711a9

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

3f8068b

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

222a638

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

57097f5

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

a42eea8

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

798f319

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

69a4972

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

run formatter

cbcb8b5

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

remove plotting

081d6ff

Merge branch 'update_advancedvi' of github.com:TuringLang/Turing.jl i…

a32a673

…nto update_advancedvi

Red-Portal marked this pull request as draft March 14, 2025 23:23

github-actions bot reviewed Mar 14, 2025

View reviewed changes

src/variational/VariationalInference.jl Outdated Show resolved Hide resolved

Red-Portal and others added 5 commits March 14, 2025 19:29

fix formatting

1bcec3e

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

fix formatting

b142832

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

fix formatting

061ec35

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

remove unused dependency

736bd3e

Merge branch 'update_advancedvi' of github.com:TuringLang/Turing.jl i…

fd434d8

…nto update_advancedvi

Red-Portal mentioned this pull request Mar 16, 2025

AdvancedVI 0.3 and compatibility with Turing.jl #2430

Open

Red-Portal requested review from yebai, mhauru and torfjelde March 16, 2025 14:08

yebai mentioned this pull request Mar 16, 2025

Interface for MAP/MLE/VI #2509

Open

Merge branch 'main' into update_advancedvi

57108ee

yebai added 2 commits March 20, 2025 21:21

Merge branch 'main' into update_advancedvi

8dc8067

Update Project.toml

297c32a

Red-Portal and others added 5 commits March 25, 2025 16:06

Merge branch 'main' of github.com:TuringLang/Turing.jl into update_ad…

3010b5e

…vancedvi

fix make some arugments of vi initializer to be optional kwargs

0c04434

Merge branch 'update_advancedvi' of github.com:TuringLang/Turing.jl i…

17a8290

…nto update_advancedvi

Merge branch 'main' into update_advancedvi

0e496c4

Update src/variational/bijectors.jl

c1533a8

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Red-Portal marked this pull request as ready for review April 20, 2025 05:26

Update Turing.jl

231d6e2

yebai requested changes Apr 21, 2025

View reviewed changes

yebai removed the request for review from torfjelde April 21, 2025 11:13

yebai assigned Red-Portal Apr 21, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update to the AdvancedVI@0.3 interface #2506

Update to the AdvancedVI@0.3 interface #2506

Red-Portal commented Mar 14, 2025 •

edited by yebai

Loading

yebai commented Mar 20, 2025

Red-Portal commented Mar 20, 2025

yebai commented Mar 20, 2025

yebai commented Apr 18, 2025

codecov bot commented Apr 21, 2025 •

edited

Loading

yebai left a comment •

edited

Loading

yebai Apr 21, 2025

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025

yebai Apr 21, 2025

Update to the AdvancedVI@0.3 interface #2506

Are you sure you want to change the base?

Update to the AdvancedVI@0.3 interface #2506

Conversation

Red-Portal commented Mar 14, 2025 • edited by yebai Loading

yebai commented Mar 20, 2025

Red-Portal commented Mar 20, 2025

yebai commented Mar 20, 2025

yebai commented Apr 18, 2025

codecov bot commented Apr 21, 2025 • edited Loading

Codecov Report

yebai left a comment • edited Loading

Choose a reason for hiding this comment

yebai Apr 21, 2025

Choose a reason for hiding this comment

yebai Apr 21, 2025 • edited Loading

Choose a reason for hiding this comment

yebai Apr 21, 2025 • edited Loading

Choose a reason for hiding this comment

yebai Apr 21, 2025 • edited Loading

Choose a reason for hiding this comment

yebai Apr 21, 2025

Choose a reason for hiding this comment

yebai Apr 21, 2025 • edited Loading

Choose a reason for hiding this comment

yebai Apr 21, 2025

Choose a reason for hiding this comment

yebai Apr 21, 2025

Choose a reason for hiding this comment

Red-Portal commented Mar 14, 2025 •

edited by yebai

Loading

codecov bot commented Apr 21, 2025 •

edited

Loading

yebai left a comment •

edited

Loading

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025 •

edited

Loading

yebai Apr 21, 2025 •

edited

Loading