Skip to content

Update to the [email protected] interface #2506

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 69 commits into from
Jun 3, 2025
Merged

Conversation

Red-Portal
Copy link
Member

@Red-Portal Red-Portal commented Mar 14, 2025

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.q_fullrank_gaussian(Random.default_rng(), m)
n_iters = 1000
Turing.Variational.vi(m, q, n_iters)

Closes #2507
Closes #2508
Closes #2430

Red-Portal and others added 12 commits March 14, 2025 19:09
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…nto update_advancedvi
@Red-Portal Red-Portal marked this pull request as draft March 14, 2025 23:23
Red-Portal and others added 5 commits March 14, 2025 19:29
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…nto update_advancedvi
@yebai
Copy link
Member

yebai commented Mar 20, 2025

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

yebai added 2 commits March 20, 2025 21:21
@Red-Portal
Copy link
Member Author

@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
Copy link
Member

yebai commented Mar 20, 2025

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.

@penelopeysm penelopeysm force-pushed the breaking branch 2 times, most recently from 16c62f2 to c2ae953 Compare May 28, 2025 20:49
Copy link
Member

@yebai yebai left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Many thanks, @Red-Portal, for the excellent work!

I only have one comment on the interface. See below.

- `state`: Collection of states used for optimization. This can be used to resume from a past call to `vi`.
- `info`: Information generated during the optimization run.
"""
function vi(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function is a thin wrapper around AdvancedVI.optimize. I'd suggest we consider #2509 (comment) for this interface to make it future-proof:

optimise(model, VI(q, n_iterations, objective; ...), ...)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given that the function vi itself is more of a legacy from the v0.2 days, I suggest we first go with this and add the new interface in the future? Whatever we decide to do, I think having vi now won't be too much of a hassle in the future. (On that matter, I am sympathetic to the proposal to unify everything into infer)

Red-Portal and others added 3 commits May 30, 2025 17:10
fix typo
Co-authored-by: Penelope Yong <[email protected]>
@Red-Portal Red-Portal requested review from penelopeysm and yebai May 30, 2025 22:47
fix forma
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@yebai yebai changed the base branch from breaking to main May 30, 2025 22:58
@yebai yebai changed the base branch from main to breaking May 30, 2025 22:59
Copy link
Member

@penelopeysm penelopeysm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just the changelog entry, please!

@Red-Portal
Copy link
Member Author

Ah crap forgot about that one. Sorry will do.

@Red-Portal Red-Portal requested a review from penelopeysm June 2, 2025 22:43
Red-Portal and others added 2 commits June 2, 2025 18:44
Copy link
Member

@penelopeysm penelopeysm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Excellent work 😄

Red-Portal and others added 2 commits June 3, 2025 12:22
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@yebai yebai merged commit ba10394 into breaking Jun 3, 2025
31 of 33 checks passed
@yebai yebai deleted the update_advancedvi branch June 3, 2025 16:19
@yebai
Copy link
Member

yebai commented Jun 3, 2025

Thanks @Red-Portal!

@yebai yebai mentioned this pull request Jun 3, 2025
15 tasks
@Red-Portal
Copy link
Member Author

@yebai @penelopeysm Thank you both for all the feedback!

penelopeysm added a commit that referenced this pull request Jun 5, 2025
* [skip ci] Bump minor version

* Clean up old code (#2574)

* Remove src/essential and deprecated function stubs

* Fix imports

* Export `@addlogprob!`

* Fix more tests

* Clean up more stuff

* Update to the [email protected] interface (#2506)

* update to match the [email protected] interface

* run formatter

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

* run formatter

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

* run formatter

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

* run formatter

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

* run formatter

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

* run formatter

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

* run formatter

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

* run formatter

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

* run formatter

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

* remove plotting

* fix formatting

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

* fix formatting

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

* fix formatting

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

* remove unused dependency

* Update Project.toml

* fix make some arugments of vi initializer to be optional kwargs

* remove tests for custom optimizers

* remove unused file

* Update src/variational/bijectors.jl

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

* Update Turing.jl

* fix remove call to `AdvancedVI.turnprogress`, which has been removed

* apply comments from @yebai

* Update src/variational/VariationalInference.jl

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

* add old interface as deprecated

* bump AdvancedVI version

* add deprecation for `meanfield`

* add `default_rng` interfaces

* add tests for variational inference

* run formatter

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

* remove "src/variational/bijectors.jl" (moved to `DynamicPPL.jl`)

* add more tests for variational inference initializer

* remove non-essential reexports, fix tests

* run formatter, rename functions

* add documentation

* fix run formatter

* fix remove debug commits

* run formatter

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

* run formatter

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

* run formatter

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

* add Variational submodule

* fix docstring style

* update docstring style

* format docstring style

* fix typo

Co-authored-by: Penelope Yong <[email protected]>

* fix use fixed seed with StableRNGs

* fix export variational families

* fix forma

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

* update changelog for advancedvi 0.4

* fix version number

* Format & add some links

* fix formatting

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

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hong Ge <[email protected]>
Co-authored-by: Penelope Yong <[email protected]>

* [no ci] fix changelog

---------

Co-authored-by: Kyurae Kim <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hong Ge <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

AdvancedVI 0.3 and compatibility with Turing.jl
5 participants