Allow access to different nutpie backends via pip-style syntax #7498

jessegrabowski · 2024-09-10T10:33:32Z

Description

Adds a pip-style syntax to the nuts_sampler argument that allows access to alternative compile backends, when relevant. This lets you get the nutpie jax backend by setting nuts_sampler='nutpie[jax]'. For backwards compatibility, nuts_sampler='nutpie' is equivalent to nuts_sampler='nutpie[numba]'.

The current PR only deals with nutpie, but we could easily extend this to include the default PyMC sampler, to compile to JAX, numba, or pytorch directly, without going through nutpie. I'm willing to do that extension in this PR if it is deemed worthwhile..

Related Issue

Closes Expose all nutpie compile backends through pm.sample #7497
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pymc--7498.org.readthedocs.build/en/7498/

jessegrabowski · 2024-09-10T10:34:50Z

I created some global variables to track available backends in nutpie, anticipating that we will eventually expand the list to include pytorch. It would be nice if these globals were defined in the nutpie library itself and we just imported them here.

ricardoV94 · 2024-09-10T11:06:36Z

The current PR only deals with nutpie, but we could easily extend this to include the default PyMC sampler, to compile to JAX, numba, or pytorch directly

I'm not sure how the other backends (specially JAX) behave with multiprocessing tbh :O Otherwise the idea sounds cool. Perhaps @aseyboldt can weigh in as he has a better picture of how we handle multiprocessing in pm.sample

jessegrabowski · 2024-09-10T11:21:59Z

Should we go work on a pmap for pytensor?

ricardoV94 · 2024-09-10T11:41:04Z

Should we go work on a pmap for pytensor?

The PyMC codebase would still need to know about it and work around differently for JAX

codecov · 2024-09-10T11:49:30Z

Codecov Report

❌ Patch coverage is 30.43478% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.27%. Comparing base (a6e6fa8) to head (eb6e216).
⚠️ Report is 15 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/sampling/mcmc.py	31.81%	15 Missing ⚠️
pymc/sampling/jax.py	0.00%	1 Missing ⚠️

❌ Your patch status has failed because the patch coverage (30.43%) is below the target coverage (50.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7498      +/-   ##
==========================================
- Coverage   90.62%   90.27%   -0.36%     
==========================================
  Files         116      116              
  Lines       18947    18983      +36     
==========================================
- Hits        17170    17136      -34     
- Misses       1777     1847      +70

Files with missing lines	Coverage Δ
pymc/sampling/jax.py	`0.00% <0.00%> (-84.10%)`	⬇️
pymc/sampling/mcmc.py	`83.39% <31.81%> (-7.61%)`	⬇️

... and 15 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2024-09-10T11:56:04Z

With numba backend we could actually do it with threads (with nogil)? Maybe worth opening an issue to investigate different backends for pymc samplers.

Also this shouldn't have to be nuts specific so for that a backend argument may make more sense. Something like this was also raised for VI recently IIRC

jessegrabowski · 2024-09-10T11:59:22Z

Should probably just be compile_kwargs like every other pymc function, but I liked my cute syntax.

ricardoV94 · 2024-09-10T12:02:40Z

Should probably just be compile_kwargs like every other pymc function, but I liked my cute syntax.

Not if we need to change how the samplers/threads are orchestrated

ricardoV94 · 2024-09-10T12:03:09Z

But compile kwargs already works anyway

jessegrabowski · 2024-09-10T12:09:17Z

But compile kwargs already works anyway

What do you mean, pm.sample already takes compile_kwargs? It's undocumented if so.

ricardoV94 · 2024-09-10T12:10:40Z

But compile kwargs already works anyway

What do you mean, pm.sample already takes compile_kwargs? It's undocumented if so.

It doesn't? Maybe I've played with global mode then

lucianopaz

@jessegrabowski, this looks very nice. I left a few suggestions though

pymc/sampling/mcmc.py

lucianopaz · 2024-09-12T10:26:52Z

tests/sampling/test_mcmc_external.py

Suggested change

@pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])

@pytest.mark.parametrize(

"nuts_sampler",

["pymc", "nutpie", "nutpie[jax]", "blackjax", "numpyro"],

)

lucianopaz · 2024-09-12T10:28:34Z

tests/sampling/test_mcmc_external.py

Suggested change

with pytest.raises(ValueError, match='Expected one of "numba" or "jax"; found "invalid"'):

with pytest.raises(

ValueError,

match='Could not parse nutpie backend. Expected one of "numba" or "jax"; found "invalid"',

):

lucianopaz · 2024-09-12T10:31:14Z

tests/sampling/test_mcmc_external.py

Suggested change

with pytest.raises(ValueError, match="Could not parse nutpie backend. Found 'nutpie[bad'"):

with pymc_model:

sample(nuts_sampler="nutpie[bad", random_seed=123, chains=2, tune=500, draws=500)

twiecki · 2024-09-14T15:47:40Z

I posted this PR with comments into GPT o1-mini:

---

Summary

This pull request introduces a pip-style syntax for specifying different Nutpie backends within the nuts_sampler argument of PyMC's pm.sample function. Users can now select a specific backend by using a syntax like nuts_sampler='nutpie[jax]', with nuts_sampler='nutpie' defaulting to nuts_sampler='nutpie[numba]' for backward compatibility. This enhancement aims to provide greater flexibility and performance optimization by allowing users to choose the most suitable backend for their computational needs.

Key Changes:

Code Enhancements:
- Backend Parsing: Introduces a helper function _extract_backend to parse and extract backend specifications from the nuts_sampler string.
- Type Definitions: Defines new type aliases for better type safety and clarity.
- Sampler Integration: Updates sampling functions to handle the new backend specifications seamlessly.
- Error Handling: Improves error messages to provide clearer guidance when invalid backend specifications are provided.
Testing Improvements:
- Expanded Test Coverage: Adds new test cases to cover various backend specifications, including valid and invalid scenarios.
- Refactored Tests: Utilizes fixtures and helper functions to reduce redundancy and enhance test maintainability.
Documentation Updates:
- Docstrings: Enhances existing docstrings to reflect the new backend syntax and usage examples.
- ReadTheDocs: Provides a preview link to the updated documentation for further review.
GitHub Discussions:
- Collaborative Feedback: Addresses feedback from reviewers, refining error handling and test cases based on suggestions.
- Future Considerations: Mentions potential future expansions, such as supporting additional backends like PyTorch.

Detailed Review

1. Code Changes

a. pymc/sampling/mcmc.py

Type Aliases:
```
ExternalNutsSampler = Literal["nutpie", "numpyro", "blackjax"]
NutsSampler = Literal["pymc"] | ExternalNutsSampler
NutpieBackend = Literal["numba", "jax"]
```
- Purpose: Enhances type safety by clearly defining allowed values for samplers and Nutpie backends.
- Suggestion: Consider placing these type aliases in a separate module if they are to be reused across different parts of the codebase.
Global Variables:
```
NUTPIE_BACKENDS = get_args(NutpieBackend)
NUTPIE_DEFAULT_BACKEND = cast(NutpieBackend, "numba")
```
- Purpose: Manages available backends and sets the default backend to numba.
- Suggestion: If the list of backends is expected to grow, consider loading them dynamically or configuring them via external settings.

Backend Extraction Function:

def extract_backend(string: str) -> NutpieBackend:
    match = re.search(r"(?<=\[)[^\]]+(?=\])", string)
    if match is None:
        return NUTPIE_DEFAULT_BACKEND
    result = cast(NutpieBackend, match.group(0))
    if result not in NUTPIE_BACKENDS:
        last_option = f"{NUTPIE_BACKENDS[-1]}"
        expected = (
            ", ".join([f'"{x}"' for x in NUTPIE_BACKENDS[:-1]]) + f' or "{last_option}"'
        )
        raise ValueError(f'Expected one of {expected}; found "{result}"')
    return result

Purpose: Parses the backend from the nuts_sampler string and validates it.
Strengths:
- Uses regex to extract backend specifications.
- Provides clear error messages when invalid backends are specified.
Suggestions:
- Regex Robustness: Ensure the regex accurately captures backends in all valid formats and handles edge cases (e.g., multiple brackets, unexpected characters).
- Error Messaging: The error message could be more descriptive by including the entire input string to aid debugging.

Sampler Integration:
```
if sampler.startswith("nutpie"):
    backend = extract_backend(sampler)
    compiled_model = nutpie.compile_pymc_model(model, backend=backend)
```
- Purpose: Integrates the backend specification into the Nutpie sampler.
- Suggestions:
  - Namespace Consistency: Ensure that all references to samplers and backends maintain consistent naming conventions.
  - Extensibility: If future backends are added, verify that the integration logic remains scalable.

b. pymc/sampling/jax.py

Type Alias Update:
```
JaxNutsSampler = Literal["numpyro", "blackjax"]
```
- Purpose: Defines allowed JAX-based samplers.
- Suggestion: Similar to ExternalNutsSampler, consider centralizing type aliases for better maintainability.
Function Signature Update:
```
def sample_jax_nuts(
    ...
    nuts_sampler: JaxNutsSampler,
    ...
) -> az.InferenceData:
```
- Purpose: Enforces that only numpyro and blackjax are accepted as JAX-based samplers.
- Suggestion: Ensure that all downstream functions and users are aware of this restriction to prevent runtime errors.

c. tests/sampling/test_mcmc_external.py

Test Refactoring:

@pytest.fixture
def pymc_model():
    with Model() as m:
        x = Normal("x", 100, 5)
        y = Data("y", [1, 2, 3, 4])
        Data("z", [100, 190, 310, 405])

        Normal("L", mu=x, sigma=0.1, observed=y)

    return m

Purpose: Defines a reusable model fixture for tests.
Strengths: Promotes DRY (Don't Repeat Yourself) principles by avoiding redundancy in test definitions.

New Test Cases:

Backend Options Test:

@pytest.mark.parametrize("backend", ["numba", "jax"], ids=["numba", "jax"])
def test_numba_backend_options(pymc_model, recwarn, backend):
    ...

Purpose: Validates that specifying different Nutpie backends works as expected.

Invalid Backend Test:

def test_invalid_nutpie_backend_raises(pymc_model):
    pytest.importorskip("nutpie")
    with pytest.raises(ValueError, match='Could not parse nutpie backend. Expected one of "numba" or "jax"; found "invalid"'):
        with pymc_model:
            sample(nuts_sampler="nutpie[invalid]", ...)

Purpose: Ensures that invalid backend specifications raise appropriate errors.
Strengths: Provides clear assertions that verify error handling mechanisms.

Test Helper Function:
```
def check_external_sampler_output(warns, idata1, idata2, sample_kwargs):
    ...
```
- Purpose: Abstracts common test assertions to enhance readability and maintainability.
- Suggestion: Ensure that this function covers all necessary assertions to comprehensively validate sampler outputs.

d. Documentation Updates

Docstring Enhancements:

nuts_sampler : str, default "pymc"
    Which NUTS implementation to run. One of ["pymc", "nutpie", "blackjax", "numpyro"]. In addition, the compilation
    backend for the chosen sampler can be set using square brackets, if available. For example, "nutpie[jax]" will
    use the JAX backend for the nutpie sampler. Currently, "nutpie[jax]" and "nutpie[numba]" are allowed.

Purpose: Clarifies the usage of the nuts_sampler argument with examples.
Strengths: Provides users with clear guidance on how to specify different backends.

Documentation Preview:
- Link: ReadTheDocs Preview
- Suggestion: Verify that all new features are accurately reflected in the documentation, including edge cases and error scenarios.

2. Testing Enhancements

Increased Coverage:
- Codecov Report: Indicates a patch coverage of 47.83% with 12 lines missing.
- Action Required: Ensure that all new code paths, especially those involving error handling and backend parsing, are adequately tested to improve coverage.
Comprehensive Test Cases:
- Valid Backend Specifications: Tests cover both numba and jax backends, ensuring they are correctly parsed and utilized.
- Invalid Backend Specifications: Tests verify that improper backend inputs are gracefully handled with meaningful error messages.
- Edge Cases: Tests include scenarios like missing brackets or unsupported backend names to ensure robustness.
Fixture Utilization:
- pymc_model Fixture: Enhances test efficiency by providing a standardized model setup for multiple tests.
- Suggestion: Consider adding more fixtures if additional setup is required for future tests, promoting scalability.

3. GitHub Discussion & Iterative Improvements

Collaborative Feedback:
- Error Handling Enhancements: Reviewer @lucianopaz suggested refining error messages and handling misformatted backend specifications, which were addressed by the author.
- Test Case Refinements: Suggestions to expand parameterizations and improve test assertions were implemented, enhancing test robustness.
Future Considerations:
- Additional Backends: The author mentioned the potential to extend support beyond Nutpie, such as integrating with PyTorch, which could be explored in future PRs.
- Multiprocessing Compatibility: Discussions around how different backends behave with multiprocessing indicate areas for further investigation to ensure seamless integration.
Commit History:
- Clarity and Relevance: Commit messages are descriptive and align with the logical progression of changes, facilitating easier codebase navigation and history tracking.
- Reversion Handling: The author promptly addressed issues raised during code review, demonstrating responsiveness and a commitment to quality.

Conclusion

This pull request effectively enhances PyMC's sampling functionality by introducing a flexible and user-friendly way to specify different Nutpie backends. The implementation is thoughtfully designed, with clear type definitions, robust error handling, and comprehensive testing to ensure reliability and maintainability.

Recommendations Before Merging:

Improve Regex Parsing:
- Robustness: Ensure that the regex used in the extract_backend function accurately captures backend specifications and handles all valid and invalid formats.
- Edge Case Handling: Consider additional tests for edge cases, such as multiple backend specifications or unexpected characters.
Enhance Test Coverage:
- Complete Coverage: Address the missing lines reported by Codecov by adding tests that cover all new code paths, especially those involving error scenarios and backend extraction.
- Performance Tests: If applicable, include tests that evaluate the performance impact of different backends to guide users in selecting the most suitable option.
Documentation Completeness:
- Usage Examples: Incorporate detailed usage examples in the documentation to demonstrate how to specify different backends effectively.
- Error Messages: Document common error messages and their resolutions to assist users in troubleshooting.
Future Backend Support:
- Scalability: Plan for the potential addition of more backends by ensuring that the current implementation is easily extensible.
- Modular Design: Consider refactoring the backend management logic into separate modules or classes to enhance modularity and separation of concerns.
Code Cleanup:
- Unused Imports: Review and remove any unused imports or variables to maintain code cleanliness.
- Consistency: Ensure that naming conventions and coding styles are consistent across the new and existing code.

Overall, this PR represents a significant improvement to PyMC's flexibility and user experience, enabling more tailored and optimized sampling strategies. With the recommended refinements, it is well-positioned for successful integration into the main codebase.

Approved with minor recommendations.

jessegrabowski · 2024-10-15T11:02:21Z

Close this as stale now that #7535 is merged? Or is there interest in this syntax still.

ricardoV94 · 2024-10-15T11:13:14Z

This syntax is a bit more ergonomic / discoverable

aseyboldt · 2024-10-16T08:32:32Z

I like the syntax :-)
But maybe we should go with jax by default? I know it is a change, but I think right now this is much more robust.

jessegrabowski · 2024-10-16T08:35:47Z

Thoughts on the best way to ask for the gradient backend in this syntax? Or nutpie[jax] should just always set both, and consider the mixed case to be "advanced"?

ricardoV94 · 2024-10-16T08:54:55Z

I like the syntax :-) But maybe we should go with jax by default? I know it is a change, but I think right now this is much more robust.

That will break people's existing code if they don't have jax installed

AlexAndorra · 2024-10-30T15:58:50Z

This looks super helpful @jessegrabowski !!

ricardoV94 · 2025-11-12T08:16:45Z

I know there was some back and forth here but I think the backend=mode kwarg would be nicer. So you can do pm.sample(backend="numba") as well as pm.sample(nuts_sampler="nutpie", backend="numba"), and pm.sample_posterior_predictive(backend="numba")

That is, the backend is somewhat orthogonal to the sampler, be it pymc or not?

Basically a less verbose shortcut for the most common use of compile_kwargs (which can still exist for more fine tweaks)

jessegrabowski · 2025-11-12T13:46:14Z

Yeah I think I agree. I've cooled on all the clever string stuff recently. I think I'll close this PR and open a fresh one with the backend argument.

lucianopaz requested changes Sep 12, 2024

View reviewed changes

ricardoV94 mentioned this pull request Oct 11, 2024

Allow for passing of backend and gradient_backend to nutpie #7535

Merged

9 tasks

jessegrabowski force-pushed the nutpie-backends branch from 763cd87 to a5b3241 Compare October 16, 2024 07:12

jessegrabowski added 5 commits November 10, 2025 21:37

Allow access to different nutpie backends via pip-style syntax

61a0895

Update pm.sample docstring

b6dfc27

Make mypy happy

3b2194d

importorskip nutpie backend test

2902cd3

Address feedback

eb6e216

jessegrabowski force-pushed the nutpie-backends branch from a5b3241 to eb6e216 Compare November 12, 2025 03:38

github-actions bot added feature request nutpie request discussion labels Nov 12, 2025

jessegrabowski closed this Nov 12, 2025

-@pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])
+@pytest.mark.parametrize(
+    "nuts_sampler",
+    ["pymc", "nutpie", "nutpie[jax]", "blackjax", "numpyro"],
+)

+    with pytest.raises(ValueError, match="Could not parse nutpie backend. Found 'nutpie[bad'"):
+        with pymc_model:
+            sample(nuts_sampler="nutpie[bad", random_seed=123, chains=2, tune=500, draws=500)

Allow access to different nutpie backends via pip-style syntax #7498

Allow access to different nutpie backends via pip-style syntax #7498

Conversation

jessegrabowski commented Sep 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

jessegrabowski commented Sep 10, 2024

Uh oh!

ricardoV94 commented Sep 10, 2024

Uh oh!

jessegrabowski commented Sep 10, 2024

Uh oh!

ricardoV94 commented Sep 10, 2024

Uh oh!

codecov bot commented Sep 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 commented Sep 10, 2024

Uh oh!

jessegrabowski commented Sep 10, 2024

Uh oh!

ricardoV94 commented Sep 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Sep 10, 2024

Uh oh!

jessegrabowski commented Sep 10, 2024

Uh oh!

ricardoV94 commented Sep 10, 2024

Uh oh!

lucianopaz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

lucianopaz Sep 12, 2024

Choose a reason for hiding this comment

Uh oh!

lucianopaz Sep 12, 2024

Choose a reason for hiding this comment

Uh oh!

lucianopaz Sep 12, 2024

Choose a reason for hiding this comment

Uh oh!

twiecki commented Sep 14, 2024

Summary

Detailed Review

1. Code Changes

2. Testing Enhancements

3. GitHub Discussion & Iterative Improvements

Conclusion

Uh oh!

jessegrabowski commented Oct 15, 2024

Uh oh!

ricardoV94 commented Oct 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aseyboldt commented Oct 16, 2024

Uh oh!

jessegrabowski commented Oct 16, 2024

Uh oh!

ricardoV94 commented Oct 16, 2024

Uh oh!

AlexAndorra commented Oct 30, 2024

Uh oh!

ricardoV94 commented Nov 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jessegrabowski commented Nov 12, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

jessegrabowski commented Sep 10, 2024 •

edited

Loading

codecov bot commented Sep 10, 2024 •

edited

Loading

ricardoV94 commented Sep 10, 2024 •

edited

Loading

ricardoV94 commented Oct 15, 2024 •

edited

Loading

ricardoV94 commented Nov 12, 2025 •

edited

Loading