Update to the [email protected] interface

[Red-Portal]

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

[yebai]

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

[Red-Portal]

@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]

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.

[mhauru]

I'm confused at the CI failure of test/mcmc/Inference.jl. It seems unrelated to this PR. I'm unable to reproduce it locally. The checks are just a tiny bit outside our allowed margin of error, but I also just find it troubling that different AD backends are giving different results at all. Maybe our seed fixing is failing. I'll rerun CI to see if it's deterministic.

[mhauru]

The reason different AD backends give different results is numerical instability in HMCDA (maybe other HMC methods too). Here's the absolute value of the difference between ForwardDiff and ReverseDiff in the first 1000 values sampled by HMCDA

The instability may in particular have something to do with the adaption phase (which is where these samples are from). Unsure about that. With a shorter adaption period the different AD backends don't diverge.

It's not a very satisfying conclusion, but I think it does mean that we can relax about different AD backends giving different results, and also about a couple of them missing our target tolerance by a small amount. I'll increase the tolerance.

I wonder if test models with nicer posteriors would not have this issue.

[mhauru]

@Red-Portal, hopefully if you merge main into this, the mcmc/Inference test failures will go away. I'll leave the merge to you though, since there are conflicts.

[Red-Portal]

@yebai I think the code is good to go on my end. How should we proceed in terms of documentation? Is this tutorial all that I need to update?

[yebai]

Is this tutorial all that I need to update?

I think so, cc @penelopeysm, who might know other related docs pages requiring update.

[penelopeysm]

The main TuringLang docs can't be edited until you've released a version of Turing that contains these fixes, so that's going to be a later PR.

First of all the API docs in docs/src/api.md will need to be edited. This is pretty much just a table that lists the identifiers that Turing re-exports, plus a link to the original docs (in this case AdvancedVI docs).

Then, this should be a minor version release, right? For Turing and DynamicPPL, we keep a breaking branch that we merge breaking changes into - the point is to aggregate breaking changes so that we don't have an overly frequent stream of minor releases.

Could you change the base branch of this PR to point to that, and edit the changelog accordingly? It might be nice to give a bit of a summary of what's new in AdvancedVI 0.4, although feel free to just point people to the API docs for the details.