Allow specification of initial state for `sample` #119

torfjelde · 2023-03-13T22:19:08Z

This seems convenient to have, e.g. for resuming sampling, running special warm-up procedures.

EDIT: Note that this is now dependent on #126

codecov · 2023-03-13T22:26:38Z

Codecov Report

Attention: 2 lines in your changes are missing coverage. Please review.

Comparison is base (4dbcb3f) 97.37% compared to head (3ed5314) 96.87%.
Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #119      +/-   ##
==========================================
- Coverage   97.37%   96.87%   -0.51%     
==========================================
  Files           8        8              
  Lines         305      320      +15     
==========================================
+ Hits          297      310      +13     
- Misses          8       10       +2

Files	Coverage Δ
src/AbstractMCMC.jl	`100.00% <ø> (ø)`
src/sample.jl	`95.87% <93.10%> (-0.78%)`	⬇️

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

src/sample.jl

devmotion

Can you add tests of all new lines?
The names init_params and initial_state seem inconsistent, can we either use init or initial in both cases?

devmotion · 2023-03-13T22:47:58Z

src/sample.jl

+                    chains = Distributed.pmap(
+                        sample_chain, pool, seeds, _init_params, _initial_state
+                    )


Does this work? I don't think pmap broadcasts its arguments?

It doesn't have to, right? _init_params will either be fill(nothing, nchains) or it will be _initial_state which should also be a vector of the correct length (I should add a check for this though)

Ah yeah, you created these arrays. The motivation for the branch here was to avoid allocating such arrays.

Yeah, am aware 👍 But the branching becomes a bit more annoying if we also have to do this for init_state, so I figured just allocating was a better. I can make it a conditional if you prefer 👍

torfjelde · 2023-03-13T22:55:42Z

Can you add tests of all new lines?

Will do (but tomorrow; am about to sleep)!

The names init_params and initial_state seem inconsistent, can we either use init or initial in both cases?

Agreed. Hmm, I guess I'll make it init_state then? Personally prefer the initial_* but given that one already exists, I guess it makes sense to add this (though this init_params isn't even offiically supported, right?)

devmotion · 2023-03-14T01:03:45Z

though this init_params isn't even offiically supported, right?

It's officially supported: It's clearly documented and used in downstream packages such as EllipticalSliceSampling, DynamicPPL, and Turing.

torfjelde · 2023-03-14T05:55:04Z

It's officially supported: It's clearly documented and used in downstream packages such as EllipticalSliceSampling, DynamicPPL, and Turing.

Yes, that I'm fully aware of! Just remembered seeing the following in the docs:

There is no "official" way for providing initial parameter values yet. However, multiple packages such as EllipticalSliceSampling.jl and AdvancedMH.jl support an init_params keyword argument for setting the initial values when sampling a single chain. To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, we decided to support init_params in the default implementations of the ensemble methods:

Which I took to mean "we haven't really made an explicit decision on how to support initial parameters, but because so many downstream packages use init_params, we stay compatible with it".

But nonetheless, I'll change it to init_state then 👍

devmotion · 2023-03-14T09:27:40Z

Generally, I think I'd prefer initial but yeah, it involves more changes 🤷 I'm also not a big fan of the name params, maybe initial_sample would be better?

Since #120 requires a breaking release anyway, we could also include more breaking changes. Otherwise we could deprecate the keyword argument, maybe something like the following could work:

function f(...; init_params=nothing, initial_params=init_params)
    if init_params !== nothing
        if initial_params !== init_params
            throw(ArgumentError("..."))
        end
        Base.deprecate("....", f)
    end
    ...
end

torfjelde · 2023-09-01T21:01:07Z

@devmotion I just came across this again; given that we decided to scrap #120 in the end, what do you think of at least merging this? Being able to specify the initial state would be quite useful.

EDIT: Just remember you're on vacation! No need to rush this:)

test/sample.jl

torfjelde · 2023-10-01T17:31:06Z

I decided to just rip the band-aid off and go with renaming init_params to initial_params right away.

If we don't, we end up in a somewhat awkward scenario where we have to pass along both init_params and initial_params, which might also just break downstream step.

devmotion · 2023-10-02T09:56:54Z

Can we merge #126 first and then rebase this PR on it? Or maybe even directly rebase the PR on #126 and adjust the base branch of the PR?

torfjelde · 2023-10-02T10:34:00Z

Most certainly 👍

torfjelde · 2023-10-02T10:42:18Z

Rebased and change base for PR.

I'll add some tests for the intial_state stuff sometime today 👍

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

isnt tested

src/sample.jl

test/sample.jl

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

devmotion · 2023-10-03T07:22:35Z

Seems the PR breaks EllipticalSliceSampling? 🤔 Interesting, even though it won't cause any breakage in practice if we tag a breaking release.

torfjelde · 2023-10-03T08:46:09Z

Could it maybe be something related to init_params -> initial_params?

But yes, I'll bump the major version so won't break anything in practice:)

devmotion · 2023-10-03T08:53:15Z

Yeah, I assume that's the reason: https://github.com/TuringLang/EllipticalSliceSampling.jl/blob/ca4babb2baba9008805bc8234a6fd182119e57dc/src/abstractmcmc.jl#L25 https://github.com/TuringLang/EllipticalSliceSampling.jl/blob/ca4babb2baba9008805bc8234a6fd182119e57dc/test/simple.jl#L40 I wonder though why other packages that currently use init_params (AdvancedMH IIRC? Turing?) are not affected by this change. Maybe they are missing tests for init_params/initial_params?

devmotion · 2023-10-03T08:59:22Z

Project.toml

@@ -9,6 +9,7 @@ version = "4.5.0"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 ConsoleProgressMonitor = "88cd18e8-d9cc-4ea6-8889-5259c0d15c8b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"


This means it could be removed from the [extras] section below. But do we actually have to depend on FillArrays? I think we should just forward the user-input or nothing, but not build any arrays explicitly?

devmotion · 2023-10-03T09:06:42Z

src/sample.jl

+    _initial_params =
+        initial_params === nothing ? FillArrays.Fill(nothing, nchains) : initial_params
+    _initial_state =
+        initial_state === nothing ? FillArrays.Fill(nothing, nchains) : initial_state


I think it would be nice to avoid these FillArrays. Maybe we could

move the function below to a callable struct (could be shared between all ensemble algorithms maybe?)

pass initial_params and initial_state to the constructor as well but only use it to define type parameters that allow us to distinguish between the four possible cases (no initial params and no initial state, only initial state, only initial params, and both initial params and state)

define the function of the callable struct depending on the type parameters, forwarding the versions with only the seed or only the seed and one additional argument to the three-argument version

Is adding in callable structs here really an improvement? 😕 I agree it's more efficient, but it seems like this will be quite a bit more complex + the efficiency doesn't really matter here, right?

A callable struct should generally be better for the compiler than a closure, shouldn't it? Regardless of whether we change or add arguments as in this PR.

Yep! But is this performance critical code? And it seems to be me that we'll need a callable struct for each scenario?

🤷 Are you sure? To me it seems one struct is sufficient - both the multithreaded and the multicore version seem to use the same inner structure, and in the serial case we could set channel = nothing. If needed we could also dispatch on the type of the algorithm to handle minor differences in the function call.

To me it seems one struct is sufficient

To clarify, I don't question whether we can have a single callable struct with different call implementations; I meant more that it seems it won't be as simple as just doing

struct SampleFunc # ... end function (f::SampleFunc)(args...) # ... end

multithreaded and the multicore version seem to use the same inner structure

But if we put initial_params and initial_state in the callable struct, then we'll need to pmap, etc. over a range containing the corresponding indices, no? Which seems like it would lead to more allocations than the current impl using Fill(nothing, nchains)?

Or am I misunderstanding what you mean here?

Just bumping this discussion:) Would be nice to get a version of this PR merged.

torfjelde · 2023-10-03T09:06:53Z

Maybe they are missing tests for init_params/initial_params?

Very likely 😕

devmotion

Let's go with the FillArrays dependency for now.

This reverts commit 8d45ff4, reversing changes made to d521815.

…tial-state"" This reverts commit 420e588.

torfjelde requested a review from devmotion March 13, 2023 22:33

github-actions bot reviewed Mar 13, 2023

View reviewed changes

src/sample.jl Outdated Show resolved Hide resolved

devmotion reviewed Mar 13, 2023

View reviewed changes

torfjelde mentioned this pull request Oct 1, 2023

Use _init_parmas for MCMCThreads and MCMCDistributed too #126

Merged

github-actions bot reviewed Oct 1, 2023

View reviewed changes

test/sample.jl Outdated Show resolved Hide resolved

test/sample.jl Outdated Show resolved Hide resolved

test/sample.jl Outdated Show resolved Hide resolved

torfjelde force-pushed the torfjelde/initial-state branch from bf70831 to 00054ef Compare October 2, 2023 10:40

torfjelde changed the base branch from master to torfjelde/init-params-fix October 2, 2023 10:41

torfjelde changed the base branch from torfjelde/init-params-fix to master October 2, 2023 14:18

torfjelde and others added 10 commits October 2, 2023 15:22

added initial_state as a kwarg

34eecab

added support for initial_state kwarg in threaded sample

fbd24ca

added support for initial_state in distributed sample

e9bee36

Update src/sample.jl

207c77d

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

removed references to _init_params and _initial_params

b67e51d

check correctness of initial states

e63c46e

renamed init_params to initial_params

5083d94

renamed references for init_params to initial_params

18ad4a5

formatting

a1c1881

initial_state missing from one mcmcsample

ca4f4b9

torfjelde force-pushed the torfjelde/initial-state branch from 00054ef to ca4f4b9 Compare October 2, 2023 14:23

fixed initial_params and initial_state for MCMCDistributed

e442880

torfjelde added 5 commits October 2, 2023 23:34

fixed typo in the initial step

fb4a1f6

replaced init_params with initial_params in tests

c6ec64e

disabled logging for large number of chains in tests where logging

74465b4

isnt tested

mroe fixes

56456c9

added tests for initial state

5d83ab4

github-actions bot reviewed Oct 2, 2023

View reviewed changes

Apply suggestions from code review

3ed5314

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

devmotion reviewed Oct 3, 2023

View reviewed changes

devmotion approved these changes Oct 24, 2023

View reviewed changes

torfjelde merged commit 8d45ff4 into master Oct 24, 2023
26 of 30 checks passed

delete-merged-branch bot deleted the torfjelde/initial-state branch October 24, 2023 15:26

torfjelde mentioned this pull request Oct 24, 2023

Maybe clean up dependence on FillArrays.jl? #131

Open

torfjelde added a commit that referenced this pull request Oct 24, 2023

Revert "Merge pull request #119 from TuringLang/torfjelde/initial-state"

420e588

This reverts commit 8d45ff4, reversing changes made to d521815.

torfjelde added a commit that referenced this pull request Oct 24, 2023

Revert "Revert "Merge pull request #119 from TuringLang/torfjelde/ini…

7ff67ae

…tial-state"" This reverts commit 420e588.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow specification of initial state for `sample` #119

Allow specification of initial state for `sample` #119

torfjelde commented Mar 13, 2023 •

edited

Loading

codecov bot commented Mar 13, 2023 •

edited

Loading

devmotion left a comment

devmotion Mar 13, 2023

torfjelde Mar 13, 2023

devmotion Mar 14, 2023

torfjelde Mar 14, 2023

torfjelde commented Mar 13, 2023

devmotion commented Mar 14, 2023

torfjelde commented Mar 14, 2023

devmotion commented Mar 14, 2023

torfjelde commented Sep 1, 2023 •

edited

Loading

torfjelde commented Oct 1, 2023

devmotion commented Oct 2, 2023

torfjelde commented Oct 2, 2023

torfjelde commented Oct 2, 2023

devmotion commented Oct 3, 2023

torfjelde commented Oct 3, 2023

devmotion commented Oct 3, 2023

devmotion Oct 3, 2023

devmotion Oct 3, 2023

torfjelde Oct 3, 2023

devmotion Oct 3, 2023 •

edited

Loading

torfjelde Oct 3, 2023

devmotion Oct 3, 2023

torfjelde Oct 3, 2023

torfjelde Oct 7, 2023

torfjelde commented Oct 3, 2023

devmotion left a comment

Allow specification of initial state for sample #119

Allow specification of initial state for sample #119

Conversation

torfjelde commented Mar 13, 2023 • edited Loading

codecov bot commented Mar 13, 2023 • edited Loading

Codecov Report

devmotion left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

torfjelde commented Mar 13, 2023

devmotion commented Mar 14, 2023

torfjelde commented Mar 14, 2023

devmotion commented Mar 14, 2023

torfjelde commented Sep 1, 2023 • edited Loading

torfjelde commented Oct 1, 2023

devmotion commented Oct 2, 2023

torfjelde commented Oct 2, 2023

torfjelde commented Oct 2, 2023

devmotion commented Oct 3, 2023

torfjelde commented Oct 3, 2023

devmotion commented Oct 3, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

devmotion Oct 3, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

torfjelde commented Oct 3, 2023

devmotion left a comment

Choose a reason for hiding this comment

Allow specification of initial state for `sample` #119

Allow specification of initial state for `sample` #119

torfjelde commented Mar 13, 2023 •

edited

Loading

codecov bot commented Mar 13, 2023 •

edited

Loading

torfjelde commented Sep 1, 2023 •

edited

Loading

devmotion Oct 3, 2023 •

edited

Loading