Add var_names argument to sample

[fonnesbeck]

Description

Allow for filtering of variables included in sampled trace via an optional var_names argument, similar to what is done for plotting.

Related Issue

• [x] Closes #7068
• [ ] Related to #

Checklist

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

Type of change

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

---

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

[fonnesbeck]

Just getting started here. Testing to see if this is the right approach.

[ricardoV94]

Apologies, I didn't see it was in draft :)

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 50.00000% with 6 lines in your changes are missing coverage. Please review.

Project coverage is 90.29%. Comparing base (aa679f3) to head (13a5d31). Report is 14 commits behind head on main.

Additional details and impacted files

[](https://app.codecov.io/gh/pymc-devs/pymc/pull/7206?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) ```diff @@ Coverage Diff @@ ## main #7206 +/- ## ========================================== - Coverage 92.29% 90.29% -2.01% ========================================== Files 100 100 Lines 16875 16896 +21 ========================================== - Hits 15575 15256 -319 - Misses 1300 1640 +340 ``` | [Files](https://app.codecov.io/gh/pymc-devs/pymc/pull/7206?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | | |---|---|---| | [pymc/backends/\_\_init\_\_.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7206?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9iYWNrZW5kcy9fX2luaXRfXy5weQ==) | `92.10% <100.00%> (+0.21%)` | :arrow_up: | | [pymc/sampling/mcmc.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7206?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zYW1wbGluZy9tY21jLnB5) | `85.96% <66.66%> (-2.04%)` | :arrow_down: | | [pymc/sampling/jax.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7206?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zYW1wbGluZy9qYXgucHk=) | `0.00% <0.00%> (-94.10%)` | :arrow_down: | ... and [6 files with indirect coverage changes](https://app.codecov.io/gh/pymc-devs/pymc/pull/7206/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

[fonnesbeck]

No worries, feedback welcome at all stages! (earlier the better, in fact)

[fonnesbeck]

Should probably enforce that all stochastic variables be included.

[ricardoV94]

Should probably enforce that all stochastic variables be included.

Maybe it's fine not to?

[fonnesbeck]

The numpyro sampler does not appear to do the right thing with the passed var names.

[ricardoV94]

The numpyro sampler does not appear to do the right thing with the passed var names.

What does it do? Checking the source code, it looks like it should if you pass just the strings?

[fonnesbeck]

Seems like at some point the names don't get converted vars.

[fonnesbeck]

(ok, fixed)

[fonnesbeck]

Not sure how code coverage drops when I've added two tests.

[ricardoV94]

Not sure how code coverage drops when I've added two tests.

Cov is flaky, not always up to date or comparing with the right commit

[ricardoV94]

Thanks @fonnesbeck !