Allow for passing of backend and gradient_backend to nutpie

[fonnesbeck]

Description

There is no means for passing backend and gradient_backend to compile_pymc_model from sample, so that JAX can be used as a backend. This just looks for these arguments in nuts_sampler_kwargs and passes them if they exist.

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

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

---

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

[ricardoV94]

Minimal implementation of #7498

[ricardoV94]

Depending on the outcome of pymc-devs/nutpie#151 it may be:

Suggested change

	backend=nuts_sampler_kwargs.pop("backend", None),
	backend=nuts_sampler_kwargs.pop("backend", "numba"),

[fonnesbeck]

As it happens, I submitted a nitpick PR to fix this. It should be None because None is handled in the function.

[ricardoV94]

Otherwise you can do something like:

compile_kwargs = {} for kwarg in ("backend", "gradient_backend"): if kwarg in nuts_sampler_kwargs: compile_kwargs[kwarg] = nuts_sampler_kwargs.pop(kwarg) compile_pymc_model(..., **compile_kwargs)

This way you don't have to guess the default of nutpie nor change anything there.

[fonnesbeck]

Will do. The arg default still needs to be changed in nutpie one way or the other, as I discuss there.

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 5 lines in your changes missing coverage. Please review.

Project coverage is 92.85%. Comparing base (c61e9cd) to head (b6de148).
Report is 67 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/sampling/mcmc.py	0.00%	5 Missing ⚠️

Additional details and impacted files

@@ Coverage Diff @@ ## main #7535 +/- ## ========================================== - Coverage 92.88% 92.85% -0.03%  ========================================== Files 105 105 Lines 17587 17591 +4 ========================================== Hits 16335 16335 - Misses 1252 1256 +4 

Files with missing lines	Coverage Δ
pymc/sampling/mcmc.py	87.23% <0.00%> (-0.74%)	⬇️