Skip to content

Optimize generated quantities replay in chain bundling#448

Open
arnavk23 wants to merge 8 commits into
TuringLang:mainfrom
arnavk23:issue-94-generated-quantities-cache
Open

Optimize generated quantities replay in chain bundling#448
arnavk23 wants to merge 8 commits into
TuringLang:mainfrom
arnavk23:issue-94-generated-quantities-cache

Conversation

@arnavk23
Copy link
Copy Markdown
Contributor

@arnavk23 arnavk23 commented Apr 15, 2026

Fixes #94 by avoiding full evaluate!! re-runs during chain bundling for HMC draws.

Changes

  • Cache generated variables on BUGSModel.
  • Add targeted replay helper to populate generated quantities from sampled parameter vectors.
  • Update JuliaBUGSMCMCChainsExt to use targeted replay instead of full model evaluation.

Impact

  • Reduces unnecessary per-sample recomputation in the HMC chain assembly path.

@arnavk23 arnavk23 marked this pull request as draft April 15, 2026 23:08
…ter space handling

- Export generated_quantities_dependency_closure from main module for CI visibility
- Fix transformed parameter handling in MCMCChains extension (HMC samples in transformed space)
- Ensure generated_quantities_dependency_closure is properly imported in Model submodule
…tion

- Use forward bijector (unconstrained->constrained) when parameters are in transformed space
- Previously was using inverse bijector which was backwards
- This fixes negative parameter values in HMC generated quantities evaluation
- Remove complex custom parameter reconstruction logic
- Delegate to evaluate_with_values!! which already handles bijectors correctly
- This function now just calls evaluate_with_values!! and returns the evaluation environment
- Ensures proper handling of both transformed (HMC) and untransformed (Gibbs) spaces
- Check if model uses auto-marginalization
- Use evaluate_with_marginalization_values!! for marginalized models
- Use evaluate_with_values!! for regular models
- This fixes BoundsError when accessing parameter vectors for marginalized models
@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 19, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.59%. Comparing base (927f61f) to head (dd069ea).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #448      +/-   ##
==========================================
+ Coverage   83.56%   83.59%   +0.02%     
==========================================
  Files          27       27              
  Lines        3949     3961      +12     
==========================================
+ Hits         3300     3311      +11     
- Misses        649      650       +1     
Flag Coverage Δ
callbacks 17.37% <40.00%> (+0.66%) ⬆️
compilation_model 33.97% <46.66%> (+0.69%) ⬆️
compilation_source_gen 34.71% <46.66%> (+0.69%) ⬆️
distributed_sampling 29.46% <46.66%> (+0.68%) ⬆️
elementary 23.47% <46.66%> (+0.72%) ⬆️
frontend 48.03% <46.66%> (+0.64%) ⬆️
gibbs 28.29% <93.33%> (+0.22%) ⬆️
graphs 18.62% <46.66%> (+0.14%) ⬆️
inference_chains 22.28% <93.33%> (+0.21%) ⬆️
inference_hmc 22.46% <93.33%> (+0.21%) ⬆️
inference_marginalization 25.29% <86.66%> (+0.07%) ⬆️
inference_mh 23.02% <86.66%> (+0.21%) ⬆️
log_density 35.79% <93.33%> (+0.22%) ⬆️
model_operations 31.05% <46.66%> (+0.70%) ⬆️
parallel_sampling 19.37% <93.33%> (+0.22%) ⬆️

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.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@arnavk23 arnavk23 marked this pull request as ready for review April 19, 2026 20:05
@arnavk23
Copy link
Copy Markdown
Contributor Author

@sunxd3 Tried working on this again.

Comment thread JuliaBUGS/ext/JuliaBUGSAdvancedHMCExt.jl Outdated
Comment thread JuliaBUGS/src/model/Model.jl Outdated
Comment thread JuliaBUGS/src/graphs.jl Outdated
Comment thread JuliaBUGS/src/JuliaBUGS.jl Outdated
Comment thread JuliaBUGS/src/model/evaluation.jl
Comment thread JuliaBUGS/src/model/evaluation.jl Outdated
…s all nodes

The docstring incorrectly stated the function only recomputes generated quantities.
In reality, it delegates to evaluate_with_values or evaluate_with_marginalization_values,
which evaluate all nodes in the model. Updated the docstring to accurately reflect this.
@github-actions
Copy link
Copy Markdown
Contributor

Benchmark Results

Julia 1.11.9 on AMD EPYC 7763 64-Core Processor (Linux)

Ratio = JuliaBUGS/Stan (lower is better for JuliaBUGS)

Model Stan Params JBUGS Params LD Ratio Grad Ratio
rats 65 65 0.35x 1.67x
pumps 12 12 0.63x 1.95x
bones 13 33 0.69x 4.46x
oxford 244 244 1.69x 4.21x
epil 303 303 0.51x 1.59x
lsat 1006 1006 0.64x 1.94x
schools 133 133 0.23x 1.56x
beetles 2 2 1.33x 4.52x
air 5 5 0.51x 1.31x

Note: Stan benchmarks use hand-optimized Stan models, not direct BUGS translations. Comparison is illustrative only.

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.

AdvancedHMCExt will use evaluate!! to rerun the model for every sample drawn, and it seems wasteful

2 participants