Skip to content

add step-setup#8351

Open
OttoVintola wants to merge 5 commits into
pymc-devs:mainfrom
OttoVintola:step-setup
Open

add step-setup#8351
OttoVintola wants to merge 5 commits into
pymc-devs:mainfrom
OttoVintola:step-setup

Conversation

@OttoVintola

@OttoVintola OttoVintola commented Jul 9, 2026

Copy link
Copy Markdown

Description

I added a function to pm.sample(...) that passes the number of draws to the step methods if they have the setup attribute. With this setup method, we can pre-allocate a vector to store TreeBatch for each sampling draw in bartrs.

Related Issue

  • Closes #
  • Related to #

Checklist

Type of change

  • New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

@read-the-docs-community

read-the-docs-community Bot commented Jul 9, 2026

Copy link
Copy Markdown

Documentation build overview

📚 pymc | 🛠️ Build #33532433 | 📁 Comparing 8138fa5 against latest (56384e5)

  🔍 Preview build  

65 files changed · + 10 added · ± 45 modified · - 10 deleted

+ Added

± Modified

- Deleted

@codecov

codecov Bot commented Jul 9, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 75.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.70%. Comparing base (56384e5) to head (02b6e19).

Files with missing lines Patch % Lines
pymc/step_methods/compound.py 60.00% 2 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #8351      +/-   ##
==========================================
- Coverage   91.73%   91.70%   -0.03%     
==========================================
  Files         128      128              
  Lines       20672    20680       +8     
==========================================
+ Hits        18963    18965       +2     
- Misses       1709     1715       +6     
Files with missing lines Coverage Δ
pymc/sampling/mcmc.py 91.94% <100.00%> (+0.03%) ⬆️
pymc/step_methods/compound.py 97.88% <60.00%> (-0.82%) ⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Comment thread pymc/sampling/mcmc.py Outdated
tune = get_default_tune_steps(step, tune)

for method in flatten_steps(step):
if hasattr(method, "setup"):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can't we add it to the baseclass so we don't need to hasattr check it?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't need it, I just added the hasattr so that pm.sample() does not throw.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see there's a set_rng... could this be passed into setup? I haven't checked when/where is it called from

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think set_rng is called per chain. If we move set_rng to setup, then the rng would be shared across chains, and vice versa if we move setup to set_rng then it would be called for each chain repeatedly. That would not be an issue in the current implementation of BART because we can check if setup has already been called. Then the method should be renamed to something like setup_chain

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah setup_chain sounds good enough, and we combine both sources of info

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.

2 participants