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 1 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%)	⬇️