Change sample_posterior_predictive API wrt to volatility

[ricardoV94]

Keep in mind this is like our second most important function in PyMC, and it was easily a footgun whenever var_names was involved or there were unobserved RVs that depended on modified data/coordinates (see #7069).

After this RV variables in the trace are never resampled by default, even if it may mean the graph is invalid (shape or semantically). A warning is issued, directing users to explicitly add variables "deemed volatile" into sample_vars (to resample them) or freeze_vars (to confirm this was intended).

Regular uses that don't touch var_names should stay the same as before. After these I think var_names becomes very intuitive (and less important as well)

Main knobs

• sample_vars — variables to force sampling. RVs present in the trace will never be resampled unless they show up here. Deterministics are still recomputed automatically based on volatility, but can be included here in case the volatility analysis wouldn't trigger them (there's a pm.do use case in the tests/docstrings). Variables that don't exist in the trace (including observeds) are always sampled, but can be added here for completeness.
• freeze_vars — variables to explicitly reuse from the trace. Should also be used for RVs in the trace that seem volatile but are intended to not be resampled. Also works for deterministics — e.g. the HSGP-style case of a training-time mean-centering that shouldn't be rederived from prediction data.
• ImplicitFreezeWarning in pymc.exceptions — emitted when a trace RV is deemed volatile but not explicitly in sample_vars or freeze_vars. This is the knob to alert users they need to make a decision, reducing silent "footguns".
• var_names is strictly output-gating now. It just determines what variables appear in the returned trace.

Other changes

• var_names / sample_vars / freeze_vars, and everything after model, are now keyword-only.
• A pm.Data replaced by a TensorConstant with a different value (e.g. via pm.do(…, make_interventions_shared=False) or freeze_dims_and_data) now participates in volatility detection (closes sample_posterior_predictive does not consider changes in "ConstantData" and "Constant coords" #6876).
• Docstring rewritten: common uses first (standard PPC, set_data predictions, freezing/forcing a deterministic, out-of-model predictions), detailed volatility rules only after.
• Some (hopefully) sane default for default var_names, when sample_vars is specified.
• Removed givens argument in the inner function (closes Remove givens_dict argument in sample posterior predictive internal function #6614)

Closes #7069
Closes #6876
Closes #6614

TODO:

• clean code
• clean tests

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #32437998 | 📁 Comparing d96c1da against latest (1b071ca)

  🔍 Preview build  

232 files changed · + 7 added · ± 225 modified

+ Added

• _modules/pymc/printing.html
• api/generated/pymc.model_table.html
• api/generated/classmethods/pymc.ode.DifferentialEquation.pullback.html
• api/generated/classmethods/pymc.ode.DifferentialEquation.pushforward.html
• api/distributions/generated/classmethods/pymc.SymbolicRandomVariable.pullback.html
• api/distributions/generated/classmethods/pymc.SymbolicRandomVariable.pushforward.html
• api/model/generated/classmethods/pymc.model.core.Model.table.html

± Modified

• genindex.html
• glossary.html
• _modules/index.html
• api/backends.html
• api/logprob.html
• api/misc.html
• api/testing.html
• contributing/developer_guide.html
• contributing/implementing_distribution.html
• _modules/pymc/data.html
• and 215 more...

[codecov]

Codecov Report

❌ Patch coverage is 97.31544% with 4 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (v6@733088a). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
pymc/sampling/forward.py	97.27%	4 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff          @@
##             v6    #8209   +/-   ##
=====================================
  Coverage      ?   91.76%           
=====================================
  Files         ?      124           
  Lines         ?    20131           
  Branches      ?        0           
=====================================
  Hits          ?    18473           
  Misses        ?     1658           
  Partials      ?        0           

Files with missing lines	Coverage Δ
pymc/exceptions.py	100.00% <100.00%> (ø)
pymc/sampling/forward.py	96.77% <97.27%> (ø)

🚀 New features to boost your workflow:

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

[ricardoV94]

There is one test failing due to pymc-devs/pytensor#2051

Verified that the patch fixes it.

Our step samplers don't eval before multiprocess splitting so numba compiles in each process... wasteful and shows the bug here.

We'll just wait but we may want to tweak the steps to eval once after compile.

Regular nuts pipeline does eval before split (in init nuts jitter)

[jessegrabowski]

General comment: the process of walking the fgraph and classifying nodes as volatile or not is exactly the function of a Feature, is there a reason why you didn't use one here?

[jessegrabowski]

Some comments.

With this refactor I think we should really think about renaming this function to something like conditional_resample then implement sample_posterior_predictive as a special case that only generates the posterior predictive. My logic is that the functionality we now offer goes well beyond just sampling the posterior predictive distribution, which is a well-defined object in the literature.

[jessegrabowski]

use sample_vars without var_names in this example?

[ricardoV94]

I think this is easier to follow, I'll add a note saying var_names wasn't needed as it would defaulted to the same

[ricardoV94]

sorry I thought this was on the section on volatility. Agree. Also removed the first example. The nudge being users shouldn't think about var_names (or at least not first). but sample_vars.

[jessegrabowski]

We might want this to be in a .. warning:: directive

[ricardoV94]

okay

[ricardoV94]

done, but it's weird with the In contrast part after ... ?

[ricardoV94]

General comment: the process of walking the fgraph and classifying nodes as volatile or not is exactly the function of a Feature, is there a reason why you didn't use one here?

A Feature is specifically about adding/keep meta information while a graph is being rewritten, we are not doing any graph rewrite here, just a snapshot analysis. I can't think of any angle where it would help

[ricardoV94]

With this refactor I think we should really think about renaming this function to something like conditional_resample then implement sample_posterior_predictive as a special case that only generates the posterior predictive. My logic is that the functionality we now offer goes well beyond just sampling the posterior predictive distribution, which is a well-defined object in the literature.

This was discussed in the original issue (and I think more internally) and rejected, Back then I was advocating for it, although I didn't have as good a name as yours (except the re- in resample, as we are not always "re-sampling").

But now I don't agree so much. What is the strict criteria for distinguishing whether this is a sample_posterior_predictive vs not? The freeze_vars is not imo. Missing vars is not (think of extra group in hierarchical models). Discarding trace variables? Again depends, maybe you are discarding because you are sampling new coordinates only, and you are just telling the function that.

So... I'm inclined to think there is no obvious set of arguments/restrictions that would constrict this function to only do "sample posterior predictive". A bit like "sample" may very well be used to get prior / prior_predictive draws (or compute deterministics), the meaning of the operation is not something you can pin-point from the computation being run.

If you agree with this, the argument becomes why we should have function with another name that has exactly the same signature/behavior? Since for most users the function does what they want, let the power users make their own alias if they want, and keep official naming simple

[jessegrabowski]

This was discussed in the original issue (and I think more internally) and rejected, Back then I was advocating for it, although I didn't have as good a name as yours (except the re- in resample, as we are not always "re-sampling").

I don't agree that the ppd is hard to pointpoint though. It's defined by the fact that we're sampling the observed distribution, conditioned on the posterior samples -- it's the base case. The advantage would be to hide complexity from the median user. That said, it is likely sufficient that pm.sample_posterior_predictive(idata) does the right thing.

If this was already discussed let's move on instead of re-litigating it ad nauseam. I accept the judgement of the collective.

[ricardoV94]

It's defined by the fact that we're sampling the observed distribution, conditioned on the posterior samples -- it's the base case.

As long as you use even one sample from the posterior it fits the mold. I don't think it must use all samples from the posterior (trivial case: dead path) or can't introduce new variables (e.g., marginalized variables that are being recovered, you are using a stochastic approximation to the posterior, conjugate posterior with closed-form solution)?

Also maybe you didn't mark your variable as observed from PyMC's standpoint (but you know this is your observation). I don't think the meaning can be decided in general.