Implement nutpie as an external sampler #7719
Open
+498
−0
Conversation
This file contains 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
- Change max_tree_depth to maxdepth to match NutPie's API - Fix coords_and_dims_for_inferencedata function call - Handle progressbar parameter correctly - Skip conversion to InferenceData since NutPie already returns one - Simplify convergence checks to handle NutPie's different structure 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
twiecki
changed the title
|
It's on my stack :) |
fonnesbeck
reviewed
Mar 22, 2025
| NUTPIE_AVAILABLE = False | ||
|
|
||
|
|
||
| class NutPie(ExternalSampler): |
There was a problem hiding this comment.
Should be Nutpie. "pie" is not capitalized in the nutpie docs, anyway.
ricardoV94
requested changes
Apr 8, 2025
| special handling in PyMC's sampling loops. | ||
| """ | ||
|
|
||
| is_external = True |
There was a problem hiding this comment.
The base class already tells us this is an ExternalSampler, no need for is_external?
| logger = logging.getLogger("pymc") | ||
|
|
||
| try: | ||
| import nutpie |
There was a problem hiding this comment.
nutpie should only be imported if the step method is used/instantiated, so as to avoid import time penalty
|
|
||
| # Create a NutPie model | ||
| logger.info("Compiling NutPie model") | ||
| nutpie_model = nutpie.compile_pymc_model( |
There was a problem hiding this comment.
I think it's better to just call nutpie.sample?
Comment on lines
+197
to
+222
| # Convert to InferenceData | ||
| if idata_kwargs is None: | ||
| idata_kwargs = {} | ||
|
|
||
| # Extract relevant variables and data for InferenceData | ||
| coords, dims = coords_and_dims_for_inferencedata(model) | ||
| constants_data = find_constants(model) | ||
| observed_data = find_observations(model) | ||
|
|
||
| # Always include sampler stats | ||
| if "include_sampler_stats" not in idata_kwargs: | ||
| idata_kwargs["include_sampler_stats"] = True | ||
|
|
||
| # NutPie already returns an InferenceData object | ||
| idata = nutpie_trace | ||
|
|
||
| # Set tuning steps attribute if possible | ||
| try: | ||
| idata.posterior.attrs["tuning_steps"] = tune | ||
| except (AttributeError, KeyError): | ||
| logger.warning("Could not set tuning_steps attribute on InferenceData") | ||
|
|
||
| # Skip compute_convergence_checks for now | ||
| # NutPie's InferenceData structure is different from PyMC's expectations | ||
|
|
||
| return idata |
There was a problem hiding this comment.
I think all of this is handled by nutpie
| Enum indicating competence level for this variable | ||
| """ | ||
| if var.dtype in continuous_types and has_grad: | ||
| return Competence.IDEAL |
|
|
||
|
|
||
| @pytest.mark.skipif(not NUTPIE_AVAILABLE, reason="NutPie not installed") | ||
| def test_nutpie_jax_backend(): |
There was a problem hiding this comment.
Reuse existing tests instead of defining new ones
|
The new test file would have to be added to the yaml file (see failing pre-commit), and I'm sure they are failing looking at it |
# for free
to join this conversation on GitHub.
Already have an account?
# to comment
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.
Summary
Test plan
tests/step_methods/test_external.pyto verify NutPie integration works🤖 Generated with Claude Code
📚 Documentation preview 📚: https://pymc--7719.org.readthedocs.build/en/7719/