Add some interface functions to support the new Gibbs sampler in Turing

[sunxd3]

The recent new Gibbs sampler provides a way forward for the Turing inference stack.

A near-to-medium-range goal has been to further reduce the glue code between Turing and inference packages (ref https://github.com/TuringLang/Turing.jl/issues/2281). The new Gibbs implementation laid a great plan to achieve this goal.

This PR is modeled after the interface of @torfjelde's recent PR. And in some aspects, it is a rehash of https://github.com/TuringLang/AbstractMCMC.jl/pull/86.

The goal of this PR is to determine and implement some necessary interface improvements, so that, when we update the inference packages up to the interface, they will more or less "just work" with the new Gibbs implementation.

As a first step, we test-flight two new functions recompute_logprob!!(rng, model, sampler, state) and getparams(state):

• recompute_logprob!!(rng, model, sampler, state) recomputes the logprob given the state
• getparams(state) extract the parameter values

Some considerations:

• This assumes a state is implemented with AbstractMCMC compatible inference packages. And a state at least stores values of parameters from the current iteration (traditionally, this is in the form of a Transition) and logprob.
• recompute_logprob!!(rng, model, sampler, state)
  • do we need rng?
  • should we make model into AbstractMCMC.LogDensityModel or just LogDensityProblem (and make inference packages depend on LogDensityProblems in the latter case)? This should allow inference packages to be independent from DynamicPPL, we can use getparams to construct a varinfo in Turing
• getparams(state)
  • What does this function return? A vector, a transition?
  • Do we need setparams?
• Do we also need some interface functions for state like getstats?

@torfjelde also says (in a Slack conversation) that the a condition(model, params) is needed, but better to be implemented by packages that defines the model, which I agree.

[sunxd3]

@yebai @devmotion @cpfiffer

[devmotion]

How is #86 related to this PR?

[torfjelde]

Hmm, it's unclear to me whether it's worth adding these methods when they have "no use" unless some notion of conditioning is also added 😕

How is https://github.com/TuringLang/AbstractMCMC.jl/pull/86 related to this PR?

getparams is probably overlapping between the two PRs, but the recompute_logprob!! method is not

[sunxd3]

I am for adding a condition interface, should we upstream this from AbstractPPL?

[yebai]

I think AbstractPPL imports AbstractMCMC, so it is also a good idea to define condition here and then reexport from AbstractPPL.

[sunxd3]

Okay, now condition and decondition are moved to AbstractMCMC from AbstractPPL.

Do we want fix here?

[sunxd3]

@devmotion @yebai @torfjelde @mhauru a penny for your thoughts?

[yebai]

Do we want fix here?

I'd keep it in DynamicPPL / AbstractPPL unless there is a reason to move here.

[torfjelde]

I'm still a bit uncertain about all of this tbh. I feel like right now we're just shoving condition and decondition (which I don't think we need for Gibbs?) into AbstractMCMC.jl to motivate the inclusion of recompute_logprob!! without much thought about whether it makes sense or not 😅

I think if this is the case, then I'm preferential to ignoring my original comment of "needing condition to motivate recompute_logprob!!", i.e. just leave it as you did originally (without condition and decondition).

[sunxd3]

I removed condition (and decondition) and use the public keyword for the new interface functions. The latter will technically change the interface, so I bumped the minor version.

I also think we should add something like AbstractState to normalize the design of state. This will introduce types for state everywhere, I am unsure of the impact. What's your thoughts on this?

[torfjelde]

I also think we should add something like AbstractState to normalize the design of state. This will introduce types for state everywhere, I am unsure of the impact. What's your thoughts on this?

Not for this PR at least:) If we want to discuss this, then we should open an issue and move discussion there.

[torfjelde]

The latter will technically change the interface, so I bumped the minor version.

It seems you've bumped the major version, not the minor version?

Also, if we're making this part of the interface, we should probably document this?

[sunxd3]

Oops, you're right.

we should probably document this?

By using the public keyword, maybe we can say "this is not official yet"? I am a little hesitate to add official documentation right now, because we don't yet have a crystal clear idea of what the interface should behave. Will add docs.

[yebai]

Some high-level comments:

• Let's introduce a setparams function to complete the getparams function.
• Let's introduce some tests to test the interface and get a more grounded view of the design.
• Think of an alternative name to recompute_logprob!!, which is a bit unintuitive in terms of what it means.

@sunxd3 please also take a careful look at

• https://github.com/TuringLang/AbstractMCMC.jl/pull/117
• https://github.com/TuringLang/AbstractMCMC.jl/pull/86
• and the DynamicHMC sampling interface design, specifically the warmup_stages and reporter arguments,

we want to push for merging these PRs and incorporate some nice ideas elsewhere in the ecosystem.