Preregister shapes of sampler stats

[michaelosthege]

As described in #6503 this adds a stats_dtypes_shape attribute to BlockedStep to replace BlockedStep.stats_dtypes.

It is implemented in a backwards-compatible manner with a deprecation warning for samplers that are not yet updated.

I also specified all stat shapes I was confident about. If someone with NUTS stat experiments could comment the remaining ones that'd be great!

Related issues

• ENH: Add BlockedStep.stats_dtypes_shapes to step signatures #6503
• Stop squeezing stats by default #6207 depends on this

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 👇

New features

• Step methods can now pre-register the shape of emitted sampler stats using the stats_dtypes_shapes class attribute. The stats_dtypes attribute is being deprecated.

[codecov]

Codecov Report

Merging #6517 (b14897e) into main (31c30dc) will decrease coverage by 71.05%.
The diff coverage is 35.89%.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6517       +/-   ##
===========================================
- Coverage   94.73%   23.69%   -71.05%     
===========================================
  Files         147      147               
  Lines       27864    27913       +49     
===========================================
- Hits        26398     6613    -19785     
- Misses       1466    21300    +19834     

Impacted Files	Coverage Δ
pymc/tests/sampling/test_mcmc.py	0.00% <0.00%> (-98.61%)	⬇️
pymc/tests/step_methods/test_compound.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/step_methods/compound.py	54.66% <52.94%> (-42.80%)	⬇️
pymc/blocking.py	87.50% <100.00%> (-8.16%)	⬇️
pymc/step_methods/hmc/hmc.py	35.71% <100.00%> (-57.15%)	⬇️
pymc/step_methods/hmc/nuts.py	95.27% <100.00%> (-2.03%)	⬇️
pymc/step_methods/metropolis.py	20.46% <100.00%> (-63.12%)	⬇️
pymc/step_methods/slicer.py	27.50% <100.00%> (-68.75%)	⬇️
pymc/tests/gp/test_gp.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/gp/test_mean.py	0.00% <0.00%> (-100.00%)	⬇️
... and 125 more

[michaelosthege]

@covertg want to take a shot a reviewing here?

[covertg]

Sure, I'll do my best! First time for pymc so obviously please take it with some grains of salt, but hope this is helpful. I don't have the experience with NUTS to confirm the shapes and dtypes so I won't comment there.

[covertg]

Main comment is regarding whether we really want to fix stats_dtypes and stats_dtypes_shapes in the constructor __new__, or whether it would be better to allow the step method to modify them later.

[covertg]

Assigning stats_dtypes and stats_dtypes_shapes here in the constructor means that they could fall out of sync later, if only one were to get modified. Is this a case we should consider? I could imagine this happening if a step method wanted to determine stat shape at initialization, for example perhaps for a stat shape that varies with the number of variables passed to the step method.

I'm not sure if that is a compelling case or not. But if it is — perhap these two attributes would be better exposed via @property with getters and setters to ensure they stay in sync?

[michaelosthege]

I agree that we should get the flexibility to have samplers initialize stats on instantiation instead of specifying them as class attributes.

I also considered properties, but this wouldn't have worked nicely with the assignment of these fields in the class definition.

I don't think that sync is a problem because we can remove the old attribute.

I will think about how a refactor to definition at initialization will look like.

[covertg]

Currently these dictionary keys are not passed to arviz when creating InferenceData objects. But if/when they are, we'll probably want a way for the user to map back from <step name> to <list of variables that were sampled by that stepper>

[michaelosthege]

Agreed. ArviZ should add such a field, but even before that gets done we can add this to mcbackend.RunMeta

[ricardoV94]

Looks good, except for a nitpick

[ricardoV94]

I would vote to use None at the PyMC level as that maps directly to the convention used in PyTensor. When we use these to create a McBackend trace then we can map None to -1.

Other backends can use None to specify unknown shape if they want.

Suggested change

	- `[4, -1]` is a matrix with 4 rows and a dynamic number of columns.
	- `[4, None]` is a matrix with 4 rows and a dynamic number of columns.