Blackjax sampler fix for breaking change / enable progress bar under parallel chain_method #7453
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
It's enough if we're compatible with the latest blackjax releases. We can raise a runtime informative error if we know the installed version of blackjax is too old to work, directing users to update it |
|
|
|
Thanks for the new release Looks like it needs to get into conda first |
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #7453 +/- ##
==========================================
- Coverage 92.20% 92.17% -0.03%
==========================================
Files 103 103
Lines 17301 17258 -43
==========================================
- Hits 15952 15908 -44
- Misses 1349 1350 +1
|
junpenglao
approved these changes
Aug 12, 2024
vandalt
pushed a commit
to vandalt/pymc
that referenced
this pull request
May 14, 2025
…parallel chain_method (pymc-devs#7453) * remove blackjax pmap warning * use gen_scan_fn * remove labels * retrigger checks * retrigger checks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
blackjax-devs/blackjax#712 changes the expected jax.lax.scan carry in
progress_bar_scan. Since pymc's external blackjax sampler directly usesprogress_bar_scanit will break whenprogressbar=True. This change switches to use a new wrapper to hide the progress bar details. In addition it enables the use of progress bars underchain_method="parallel".I think any breaking issues can be handled by restricting blackjax version numbers. However, I'm not sure how to properly do that?
And of course for now tests are expected to fail until the changes show in a blackjax release.
PRs that are dependencies:
blackjax-devs/blackjax#712
blackjax-devs/blackjax#716
Checklist
Type of change
📚 Documentation preview 📚: https://pymc--7453.org.readthedocs.build/en/7453/