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

sunxd3 · 2024-07-12T15:43:38Z

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 TuringLang/Turing.jl#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 #86.

(the explanation here is outdated, please refer to #144 (comment))

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?~~

~~Tor 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 · 2024-07-12T15:45:33Z

@yebai @devmotion @cpfiffer

devmotion · 2024-07-14T21:23:19Z

How is #86 related to this PR?

torfjelde · 2024-07-15T07:00:47Z

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 #86 related to this PR?

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

sunxd3 · 2024-07-16T07:59:53Z

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

yebai · 2024-07-16T14:00:04Z

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

sunxd3 · 2024-07-18T10:18:38Z

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

Do we want fix here?

sunxd3 · 2024-07-19T10:17:27Z

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

yebai · 2024-07-19T11:44:48Z

Do we want fix here?

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

src/AbstractMCMC.jl

torfjelde · 2024-07-19T19:27:20Z

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 · 2024-07-22T07:42:28Z

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 · 2024-07-22T08:58:43Z

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 · 2024-07-22T08:59:40Z

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 · 2024-07-22T09:27:59Z

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 · 2024-07-23T19:05:39Z

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

Addition of step_warmup #117
Add getparameters and setparameters!! #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.

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

codecov · 2024-09-22T20:25:48Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.54%. Comparing base (2a77f53) to head (3ed5cb3).
Report is 5 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #144      +/-   ##
==========================================
+ Coverage   97.19%   97.54%   +0.34%     
==========================================
  Files           8        8              
  Lines         321      326       +5     
==========================================
+ Hits          312      318       +6     
+ Misses          9        8       -1

Flag	Coverage Δ
	`97.54% <ø> (?)`

Flags with carried forward coverage won't be shown. Click here to find out more.

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

sunxd3 · 2024-10-01T20:14:02Z

@yebai @torfjelde this has been updated, many issues in the previous version has been corrected (thanks for the discussion and code review). I also added another notes in the folder design_notes, comments are welcomed. Can you give it another read?

sunxd3 · 2024-10-01T20:38:43Z

In its current form, no interface change is made to AbstractMCMC, all the interface functions are from other packages.

sunxd3 · 2024-10-03T12:30:06Z

the test error seems to be Julia 1.6-only related

yebai

I left some comments below. Some other high-level comments:

unify MCMCState with MCMCTransition by introducing an AbstractMCMCState type
replace Base.vec(state) with a special state AbstractMCMC.VectorMCMCState{T}<:AbstractVector{T}, which supports getindex and setindex. All sampler packages should explicitly implement a VectorMCMCState(state) type conversion funciton.
is MCMCTransition replaceable with VectorMCMCState?

design_notes/on_gibbs_implementation.md

docs/src/state_interface.md

yebai · 2024-10-03T14:41:25Z

docs/src/state_interface.md

+This function takes the state and returns a vector of the parameter values stored in the state.
+
+```julia
+state = StateType(state::StateType, logp)


Suggested change

state = StateType(state::StateType, logp)

state = StateType(state::StateType, logdensity=logp)

yebai · 2024-10-03T14:43:01Z

docs/src/state_interface.md

+This function takes an existing `state` and a log probability value `logp`, and returns a new state of the same type with the updated log probability.
+
+These functions provide a minimal interface to interact with the `state` datatype, which a sampler package can optionally implement.
+The interface facilitates the implementation of "meta-algorithms" that combine different samplers.


Suggested change

The interface facilitates the implementation of "meta-algorithms" that combine different samplers.

The interface facilitates the implementation of "high-order" MCMC sampling algorithms like Gibbs.

docs/src/state_interface.md

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

sunxd3 · 2024-10-03T15:50:20Z

related reply from @devmotion TuringLang/Turing.jl#2304 (comment)

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

torfjelde

I'm very much in favour of using explicit functions for the interface and not overloading methods (in ways that the original methods were not intended for) 😕

torfjelde · 2024-10-04T09:54:36Z

design_notes/on_gibbs_implementation.md

+
+Here, some alternative functions that achieve the same functionality as `getparams` and `recompute_logp!!` are proposed, but without introducing new interface functions.
+
+For `getparams`, we can use `Base.vec`. It is a `Base` function, so there's no need to export anything from `AbstractMCMC`. Since `getparams` should return a vector, using `vec` makes sense. The concern is that, officially, `Base.vec` is defined for `AbstractArray`, so it remains a question whether we should only introduce `vec` in the absence of other `AbstractArray` interfaces.


I'd much prefer an explicit method in AbstractMCMC (uncertain if we want to export it 🤷 but probably make it public). Anyone implementing this interface already has AbstractMCMC loaded, so really doesn't cost anything + avoids misuse of Base.

I can resonate, issue with public is they still count as public interface, unsure if we need to make minor release

torfjelde · 2024-10-04T09:56:40Z

design_notes/on_gibbs_implementation.md

+For `recompute_logp!!`, we could overload `LogDensityProblems.logdensity(logdensity_model::AbstractMCMC.LogDensityModel, state::State; recompute_logp=true)` to compute the log probability. If `recompute_logp` is `true`, it should recompute the log probability of the state. Otherwise, it could use the log probability stored in the state. To allow updating the log probability stored in the state, samplers should define outer constructor for their state types `StateType(state::StateType, logdensity=logp)` that takes an existing `state` and a log probability value `logp`, and returns a new state of the same type with the updated log probability.
+
+While overloading `LogDensityProblems.logdensity` to take a state object instead of a vector for the second argument somewhat deviates from the interface in `LogDensityProblems`, it provides a clean and extensible solution for handling log probability recomputation within the existing interface.


But here we're introducing kwargs, etc. which is really not a part of the LogDensityProblems.logdensity interface. It would also mean we would have to depend on LogDensityProblems.jl, which we're currently not doing (AFIAK).

Why would we do this vs. just using recompute_logp!! for this?

mainly to not make any changes to the interface, so these are just "recommendations"

I think AbstractMCMC depends on LogDensityProblems

torfjelde · 2024-10-04T09:59:48Z

design_notes/on_gibbs_implementation.md

+
+## Proposed Interface
+
+The two functions `getparams` and `recompute_logp!!` form a minimal interface to support the `Gibbs` implementation. However, there are concerns about introducing them directly into `AbstractMCMC`. The main reason is that `AbstractMCMC` is a root dependency of the `Turing` packages, so we want to be very careful with new releases.


The main reason is that AbstractMCMC is a root dependency of the Turing packages, so we want to be very careful with new releases.

Fair, but if we now make a release where we assume that certain functionality is overloaded, then that seems strictly worse, no?

torfjelde · 2024-10-04T10:21:33Z

I do think the entire process of this would be quite a bit less painful if we did the following (I believe I've mentioned this before; if not, I apologize):

Improve Add getparameters and setparameters!! #86 to a finalized form . This is useful, not just for Gibbs sampling.
Make a separate package, e.g. AbstractMCMCGibbs.jl, which implements the Gibbs-only stuff, e.g. recompute_logprob!! and the sampler mapping stuff.

This is how we're doing it with MCMCTempering.jl, i.e. keep it as a separate package and slowly move pieces to AbstractMCMC.jl if it seems suitable. My problem, as stated before, is that the current Gibbs impls we're working with are really not good enough as I think is evident by a) issues that we've encountered with my Turing.jl-impl in TuringLang/Turing.jl#2328 (comment), and b) the amount of iterating you've done in this PR. This shit is complicated 😬 And I imagine it's really annoying iterating on this back and forth but without actually getting stuff merged..

So, I think a separate package would just make this entire process much easier @sunxd3 ; then we can iterate much faster on ideas (just make breaking releases), and then we can just upstream changes as we finalize things there + we can even inform about this in the official AbstractMCMC.jl docs and then people can easily support this via extensions.

torfjelde · 2024-10-04T15:06:10Z

#85 (comment)

sunxd3 added 2 commits July 12, 2024 09:26

very incomplete draft

dcf1da9

update getparams

cdaa663

Upstream condition and decondition from AbstractPPL

57275f5

torfjelde reviewed Jul 19, 2024

View reviewed changes

src/AbstractMCMC.jl Outdated Show resolved Hide resolved

src/AbstractMCMC.jl Outdated Show resolved Hide resolved

sunxd3 added 3 commits July 22, 2024 08:25

remove condition and decondition

26027ea

add Compat to make new interface functions public

6ebab49

bump minor version

e1099f9

bump minor version instead

95d781b

unfinished gibbs example

f05f293

sunxd3 mentioned this pull request Aug 12, 2024

Slice sampling as a Gibbs sampler TuringLang/Turing.jl#2300

Open

sunxd3 self-assigned this Aug 12, 2024

sunxd3 and others added 5 commits August 14, 2024 18:35

some updates

590d37f

more progress; still need to deal with w being on simplex

3afc232

bit of format

55dbab5

results is wrong

67ff8e8

Apply suggestions from code review

f758a4c

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

sunxd3 added 2 commits September 22, 2024 20:10

update doc -- need proofread

b798b2e

fix 1.6 struct field splatting compat issue

3ed5cb3

torfjelde mentioned this pull request Sep 23, 2024

Revive the package TuringLang/MCMCTempering.jl#159

Open

6 tasks

sunxd3 added 9 commits September 27, 2024 17:01

update code and doc

6fde198

relax test error

c7f577d

rename gibbs markdown file

8f11a15

change title

48a160d

update code and note

8d74889

fix doc example

bceb510

try to fix doc example error

c177271

fix doc deps

bdba893

fix more doc example error

e7e2870

minor update

80df187

sunxd3 marked this pull request as ready for review October 1, 2024 20:19

yebai reviewed Oct 3, 2024

View reviewed changes

Apply suggestions from code review

076e431

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

sunxd3 and others added 2 commits October 3, 2024 23:50

Update docs/src/state_interface.md

4293868

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Update docs/src/state_interface.md

1cee0ab

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

torfjelde requested changes Oct 4, 2024

View reviewed changes

sunxd3 marked this pull request as draft October 14, 2024 12:12

sunxd3 closed this Oct 15, 2024

sunxd3 mentioned this pull request Oct 15, 2024

Add a note on future interface for Gibbs #148

Open

yebai deleted the sunxd/interface_for_gibbs branch October 22, 2024 20:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

sunxd3 commented Jul 12, 2024 •

edited

Loading

sunxd3 commented Jul 12, 2024

devmotion commented Jul 14, 2024

torfjelde commented Jul 15, 2024

sunxd3 commented Jul 16, 2024

yebai commented Jul 16, 2024

sunxd3 commented Jul 18, 2024

sunxd3 commented Jul 19, 2024

yebai commented Jul 19, 2024 •

edited

Loading

torfjelde commented Jul 19, 2024 •

edited

Loading

sunxd3 commented Jul 22, 2024

torfjelde commented Jul 22, 2024

torfjelde commented Jul 22, 2024

sunxd3 commented Jul 22, 2024 •

edited

Loading

yebai commented Jul 23, 2024 •

edited

Loading

codecov bot commented Sep 22, 2024

sunxd3 commented Oct 1, 2024

sunxd3 commented Oct 1, 2024

sunxd3 commented Oct 3, 2024

yebai left a comment •

edited

Loading

yebai Oct 3, 2024

yebai Oct 3, 2024

sunxd3 commented Oct 3, 2024

torfjelde left a comment

torfjelde Oct 4, 2024

sunxd3 Oct 4, 2024

torfjelde Oct 4, 2024

sunxd3 Oct 4, 2024

torfjelde Oct 4, 2024

torfjelde commented Oct 4, 2024 •

edited

Loading

torfjelde commented Oct 4, 2024

	state = StateType(state::StateType, logp)
	state = StateType(state::StateType, logdensity=logp)

	The interface facilitates the implementation of "meta-algorithms" that combine different samplers.
	The interface facilitates the implementation of "high-order" MCMC sampling algorithms like Gibbs.


		Here, some alternative functions that achieve the same functionality as `getparams` and `recompute_logp!!` are proposed, but without introducing new interface functions.

		For `getparams`, we can use `Base.vec`. It is a `Base` function, so there's no need to export anything from `AbstractMCMC`. Since `getparams` should return a vector, using `vec` makes sense. The concern is that, officially, `Base.vec` is defined for `AbstractArray`, so it remains a question whether we should only introduce `vec` in the absence of other `AbstractArray` interfaces.

		For `recompute_logp!!`, we could overload `LogDensityProblems.logdensity(logdensity_model::AbstractMCMC.LogDensityModel, state::State; recompute_logp=true)` to compute the log probability. If `recompute_logp` is `true`, it should recompute the log probability of the state. Otherwise, it could use the log probability stored in the state. To allow updating the log probability stored in the state, samplers should define outer constructor for their state types `StateType(state::StateType, logdensity=logp)` that takes an existing `state` and a log probability value `logp`, and returns a new state of the same type with the updated log probability.

		While overloading `LogDensityProblems.logdensity` to take a state object instead of a vector for the second argument somewhat deviates from the interface in `LogDensityProblems`, it provides a clean and extensible solution for handling log probability recomputation within the existing interface.


		## Proposed Interface

		The two functions `getparams` and `recompute_logp!!` form a minimal interface to support the `Gibbs` implementation. However, there are concerns about introducing them directly into `AbstractMCMC`. The main reason is that `AbstractMCMC` is a root dependency of the `Turing` packages, so we want to be very careful with new releases.

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

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

Conversation

sunxd3 commented Jul 12, 2024 • edited Loading

sunxd3 commented Jul 12, 2024

devmotion commented Jul 14, 2024

torfjelde commented Jul 15, 2024

sunxd3 commented Jul 16, 2024

yebai commented Jul 16, 2024

sunxd3 commented Jul 18, 2024

sunxd3 commented Jul 19, 2024

yebai commented Jul 19, 2024 • edited Loading

torfjelde commented Jul 19, 2024 • edited Loading

sunxd3 commented Jul 22, 2024

torfjelde commented Jul 22, 2024

torfjelde commented Jul 22, 2024

sunxd3 commented Jul 22, 2024 • edited Loading

yebai commented Jul 23, 2024 • edited Loading

codecov bot commented Sep 22, 2024

Codecov Report

sunxd3 commented Oct 1, 2024

sunxd3 commented Oct 1, 2024

sunxd3 commented Oct 3, 2024

yebai left a comment • edited Loading

Choose a reason for hiding this comment

yebai Oct 3, 2024

Choose a reason for hiding this comment

yebai Oct 3, 2024

Choose a reason for hiding this comment

sunxd3 commented Oct 3, 2024

torfjelde left a comment

Choose a reason for hiding this comment

torfjelde Oct 4, 2024

Choose a reason for hiding this comment

sunxd3 Oct 4, 2024

Choose a reason for hiding this comment

torfjelde Oct 4, 2024

Choose a reason for hiding this comment

sunxd3 Oct 4, 2024

Choose a reason for hiding this comment

torfjelde Oct 4, 2024

Choose a reason for hiding this comment

torfjelde commented Oct 4, 2024 • edited Loading

torfjelde commented Oct 4, 2024

sunxd3 commented Jul 12, 2024 •

edited

Loading

yebai commented Jul 19, 2024 •

edited

Loading

torfjelde commented Jul 19, 2024 •

edited

Loading

sunxd3 commented Jul 22, 2024 •

edited

Loading

yebai commented Jul 23, 2024 •

edited

Loading

yebai left a comment •

edited

Loading

torfjelde commented Oct 4, 2024 •

edited

Loading