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

• Closes Allow storing subset of variable in pm.sample via a var_names kwarg #7068
• 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

• 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]

Thanks @fonnesbeck

Left a small comment. Also I think the jax samplers kind of allow this functionality, so would only need to forward to the sample_external_nuts code path?

[ricardoV94]

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

[codecov]

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

@@            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	Coverage Δ
pymc/backends/__init__.py	92.10% <100.00%> (+0.21%)	⬆️
pymc/sampling/mcmc.py	85.96% <66.66%> (-2.04%)	⬇️
pymc/sampling/jax.py	0.00% <0.00%> (-94.10%)	⬇️

... and 6 files with indirect coverage changes

[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]

Add a warning about var_names not beeing used by nutpie like we have for some other arguments above?

[ricardoV94]

Also, @aseyboldt how hard/reasonable is it to support this in nutpie?

[fonnesbeck]

Perhaps it could be filtered in nutpie.sample._trace_to_arviz?

[ricardoV94]

Ideally we want to filter during sampling already since RAM is usually the issue, not disk-space?

[fonnesbeck]

I agree, but there are no obvious hooks into the nutpie compiled model. It would require some changes on the nutpie side, by the looks of it.

[ricardoV94]

My top comment was to add a warning like these:

pymc/pymc/sampling/mcmc.py

Lines 283 to 294 in abe7bc9

	if initvals is not None:
	warnings.warn(
	"`initvals` are currently not passed to nutpie sampler. "
	"Use `init_mean` kwarg following nutpie specification instead.",
	UserWarning,
	)
	if idata_kwargs is not None:
	warnings.warn(
	"`idata_kwargs` are currently ignored by the nutpie sampler",
	UserWarning,
	)

Not to try to monkey-patch nutpie from the outside

[ricardoV94]

Once/if nutpie has similar functionality we can forward it from pymc?

[fonnesbeck]

Let's go with the warning for now, and create an issue on nutpie for a solution.

[aseyboldt]

Shouldn't be too hard. Nutpie uses a numba function to compute all values that should appear in the trace (including the deterministics and transformed values). We should be able to just export a subset (code is around here: https://github.com/pymc-devs/nutpie/blob/main/python/nutpie/compile_pymc.py#L387)

[ricardoV94]

Thanks @fonnesbeck !