Skip to content
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

Fix for generated_quantities #534

Merged
merged 31 commits into from
Sep 9, 2023
Merged

Conversation

torfjelde
Copy link
Member

@torfjelde torfjelde commented Sep 6, 2023

As mentioned elsewhere, now that there is no longer such an easy mapping between the strings in the a resulting Chains to the variables in a given VarInfo, some methods, e.g. generated_quantities don't work so well (TuringLang/Turing.jl#2077).

This PR is one attempt at fixing this. It's not particularly satisfying but I struggle to see another way around which wouldn't require substantial changes to MCMCChains.jl.

The idea is basically as follows:

  • In the chains.info field we add a dict-like mapping from VarName to Symbol.
  • This can then be used to index into the chains using a VarName by simply mapping the VarName to the corresponding Symbol, and performing indexing as usual.
  • Add a nested_getindex and nested_setindex! which "tries a bit harder" to well, getindex and setindex!. When I say "tries a bit harder" I mean that if we can't quite locate the corresponding varname, we see if we can identify a varname that subsumes it (a parent_varname) and then use the "remainder"-lens to interact with the reconstructed object.

The result: we preserve the structure-information present in VarName in Chains, and using nested_setindex! we can then correctly updated more structural variables, e.g. `Cholesky** variables.

The downside: this is completely reliant on the Chains being constructed with the info field having this varname-to-symbol mapping. Luckily, we can easily enforce this in chains returned from Turing.jl.

Thoughts?

Fixes #515

ext/DynamicPPLMCMCChainsExt.jl Outdated Show resolved Hide resolved
ext/DynamicPPLMCMCChainsExt.jl Outdated Show resolved Hide resolved
src/utils.jl Outdated Show resolved Hide resolved
src/varinfo.jl Outdated Show resolved Hide resolved
test/model.jl Outdated Show resolved Hide resolved
torfjelde and others added 2 commits September 6, 2023 08:14
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
src/varinfo.jl Outdated Show resolved Hide resolved
test/model.jl Outdated Show resolved Hide resolved
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@yebai yebai requested a review from sunxd3 September 6, 2023 10:47
@yebai
Copy link
Member

yebai commented Sep 6, 2023

I'm happy with the general idea. We could also store the mapping inside DynamicPPL.Model so that the mapping functionality is independent of MCMCChains. This might be helpful in the future when we support InferenceObjects.

@torfjelde
Copy link
Member Author

We could also store the mapping inside DynamicPPL.Model so that the mapping functionality is independent of MCMCChains.

I don't quite see how that would work; can you elaborate?

torfjelde and others added 3 commits September 6, 2023 12:45
src/utils.jl Outdated Show resolved Hide resolved
src/utils.jl Outdated Show resolved Hide resolved
src/utils.jl Outdated Show resolved Hide resolved
src/utils.jl Outdated Show resolved Hide resolved
torfjelde and others added 2 commits September 6, 2023 13:42
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
src/varinfo.jl Outdated Show resolved Hide resolved
@yebai
Copy link
Member

yebai commented Sep 6, 2023

I mean we could store chain.info.varname_to_symbol in DynamicPPL.model instead.

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@torfjelde
Copy link
Member Author

I mean we could store chain.info.varname_to_symbol in DynamicPPL.model instead.

Hmm, I don't think I like that. IMO the Model shouldn't hold any state or "model-runtime"-dependent information; which variables are available can easily change between executions / instantiations. We really don't have a good way of getting the varnames without just executing the model, and so it seems a bit strange to me to then attach this information directly to the model.

@yebai
Copy link
Member

yebai commented Sep 6, 2023

Hmm, not entirely certain what to do about this. Do we introduce Requires.jl? Do we simply not run these particular tests on lower versions?

Requires has been widely adopted for handling weak dependencies for pre-1.9 releases, so I think it is safe to take this approach.

src/DynamicPPL.jl Outdated Show resolved Hide resolved
torfjelde and others added 3 commits September 6, 2023 18:49
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@codecov
Copy link

codecov bot commented Sep 6, 2023

Codecov Report

Patch coverage: 87.50% and project coverage change: +0.07% 🎉

Comparison is base (52cd7f9) 81.24% compared to head (46e5a94) 81.32%.

Additional details and impacted files
@@            Coverage Diff             @@
##           master     #534      +/-   ##
==========================================
+ Coverage   81.24%   81.32%   +0.07%     
==========================================
  Files          24       25       +1     
  Lines        2932     2982      +50     
==========================================
+ Hits         2382     2425      +43     
- Misses        550      557       +7     
Files Changed Coverage Δ
src/DynamicPPL.jl 33.33% <0.00%> (-66.67%) ⬇️
src/chains.jl 0.00% <0.00%> (ø)
src/varinfo.jl 92.27% <87.50%> (-0.27%) ⬇️
ext/DynamicPPLMCMCChainsExt.jl 100.00% <100.00%> (ø)
src/utils.jl 81.49% <100.00%> (+0.44%) ⬆️

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

src/DynamicPPL.jl Outdated Show resolved Hide resolved
Co-authored-by: David Widmann <[email protected]>
@torfjelde
Copy link
Member Author

Should we maybe wait until ConstructionBase.jl gets its version bump? Seems like the PR isn't far from going through.

so we can make use of this also for Turing.predict, etc.
@torfjelde
Copy link
Member Author

I think it might be worth having another look at this PR @yebai @sunxd3 @devmotion

I added a couple of methods + moved the new functionality in to setval_and_resample! so it's more easily re-usable + immediately becomes useful for existing methods of similar nature to generated_quantities, e.g. Turing.predict

ext/DynamicPPLMCMCChainsExt.jl Outdated Show resolved Hide resolved
ext/DynamicPPLMCMCChainsExt.jl Outdated Show resolved Hide resolved
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>
src/chains.jl Outdated Show resolved Hide resolved
@sunxd3
Copy link
Member

sunxd3 commented Sep 7, 2023

Make sense to me!

@torfjelde
Copy link
Member Author

This should now be good to go once the tests pass:)

@yebai yebai enabled auto-merge September 8, 2023 20:19
@yebai yebai added this pull request to the merge queue Sep 8, 2023
Merged via the queue into master with commit ffe9272 Sep 9, 2023
12 of 13 checks passed
@yebai yebai deleted the torfjelde/nested-get-and-setindxe branch September 9, 2023 01:45
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.

MCMCChains for Julia < 1.9
4 participants