Check shared variable values to determine volatility in posterior predictive sampling

[lucianopaz]

What is this PR about?

Closes #6047.

This PR did two things:

1. It makes the mutable dimension SharedVariable length have the name of the corresponding dimension.
2. Adds two arguments to compile_forward_sampling_function: constant_data and constant_coords.

These two arguments allow it to determine if a SharedVariable has changed after between a call to sample and sample_posterior_predictive. We have to note that constant_data is only knowable if sample_posterior_predictive is called supplying an InferenceData object, and constant_coords are only knowable if sample_posterior_predictive is called supplying an InferenceData or an xarray.Dataset object.

The way sample_posterior_predictive is able to determine if a data container changed is by doing the following. It first checks the InferenceData.constant_data group to find the values of data containers at inference time, and passes those into the constant_data argument of compile_forward_sampling_function. When a SharedVariable is found while walking the graph, it looks up the entry in constant_data, it it finds an entry, it checks whether the values in the dictionary match the values of the SharedVariable at run time. If they match, the SharedVariable is deemed not volatile, if they don't match, the SharedVariable is considered volatile.
To check if a dimension's coordinates changed, sample_posterior_predictive compares the model's coordinates to those found in the supplied trace (must be an InferenceData or xarray.Dataset to have this information). If the coordinates did not change, then the dimension name is added to the constant_coords set. Then, when compile_forward_sampling_function finds the dimension length's shared variable, it tries to see if its name is in constant_coords. If it is, then the dimension is not deemed volatile, if it isn't it is considered volatile.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• ...

Bugfixes / New features

• sample_posterior_predictive, when supplied with an InferenceData object, properly identifies if a MutableData or mutable dimension has changed between a call to pymc.sample and pymc.sample_posterior_predictive. If they have, then the descendant random variables are resampled, if they have not changed, then the descendant random variables are taken from the InferenceData.posterior.

Docs / Maintenance

• ...

[codecov]

Codecov Report

Merging #6147 (b12c189) into main (eff1cf2) will increase coverage by 47.10%.
The diff coverage is n/a.

❗ Current head b12c189 differs from pull request most recent head a9f6ab5. Consider uploading reports for the commit a9f6ab5 to get more accurate results

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6147       +/-   ##
===========================================
+ Coverage   40.84%   87.94%   +47.10%     
===========================================
  Files          91       91               
  Lines       20713    20707        -6     
===========================================
+ Hits         8460    18211     +9751     
+ Misses      12253     2496     -9757     

Impacted Files	Coverage Δ
pymc/tests/distributions/test_bound.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/distributions/test_censored.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/distributions/test_truncated.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/distributions/test_simulator.py	0.00% <0.00%> (-99.51%)	⬇️
pymc/sampling_jax.py	0.00% <0.00%> (-97.19%)	⬇️
pymc/distributions/truncated.py	34.96% <0.00%> (-64.34%)	⬇️
pymc/distributions/bound.py	45.63% <0.00%> (-54.37%)	⬇️
pymc/distributions/simulator.py	34.04% <0.00%> (-53.20%)	⬇️
pymc/distributions/censored.py	92.50% <0.00%> (ø)
pymc/parallel_sampling.py	85.80% <0.00%> (+0.33%)	⬆️
... and 69 more

[lucianopaz]

I'm not sure why precommit just dies after checking for links.

[ricardoV94]

Neat one @lucianopaz!