Refactor backend split: unify pymc and numpyro simulation loops - #60
Conversation
Consolidate the per-backend simulation loops into a single shared loop, move pymc prior/posterior predictive into the adapter, apply a strategy pattern for rank calculation, and clean up typing, logging, and ruff issues. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #60 +/- ##
==========================================
- Coverage 96.69% 95.87% -0.82%
==========================================
Files 4 7 +3
Lines 545 631 +86
==========================================
+ Hits 527 605 +78
- Misses 18 26 +8 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
aloctavodia
left a comment
aloctavodia
left a comment
There was a problem hiding this comment.
This is great! I am leaving a couple of comments for future reference.
| thinned_idata = self.trace.isel(draw=draw_indices) | ||
| posterior = extract(thinned_idata, group="posterior", keep_dataset=True) |
There was a problem hiding this comment.
Not really related to this PR, but leaving the comment before I forget. Do we want to do this, or do we want to use https://python.arviz.org/projects/stats/en/stable/api/generated/arviz_stats.thin.html, which by default thins based on the ESS?
There was a problem hiding this comment.
if we are to do this, we should do this for all other instances of sampling?
|
|
||
| return posterior, posterior_pred | ||
|
|
||
| def compute_single_rank(self, transform, name, posterior, simulation_idx, ref_params): |
There was a problem hiding this comment.
Again, for the future. We may want to do this directly using xarray
4 participants
In main, SBC class has three coupled concepts: the SBC loop itself, PyMC-specific sampling and model manipulation details and Numpyro-specifics.
This PR:
decouples SBC loop from PyMC from Numpyro. As a result, now there's a single simulation loop, backend-agnostic.