fix: batched mcmc reshaping

[janfb]

when the num_chains was higher or not a multiple of the number of samples to be generated with MCMC sample_batched, the samples were reshaped incorrectly: The remaining samples were pushed into the last dimension here:

sbi/sbi/inference/posteriors/mcmc_posterior.py

Line 467 in b275448

return samples.reshape((batch_size, *sample_shape, -1)).permute( # type: ignore

This is a fix where we first collect all the generated samples from the chains (possibly more than needed), then select as many as we need and then reshape into the desired (*sample_shape, batch_size, input_shape).

I also added tests that cover all the cases, which makes them quite slow.

[janfb]

@gmoss13 I had to change one thing in your reshape-permute magic to make it work for len(sample_shape) > 1. Can you please double check that this does not mess up the collection of samples from the chains?

[codecov]

Codecov Report

Attention: Patch coverage is 91.66667% with 1 line in your changes missing coverage. Please review.

Project coverage is 75.86%. Comparing base (ba19688) to head (7e76605).
Report is 9 commits behind head on main.

Files	Patch %	Lines
sbi/inference/posteriors/mcmc_posterior.py	91.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1210      +/-   ##
==========================================
- Coverage   84.55%   75.86%   -8.70%     
==========================================
  Files          96       97       +1     
  Lines        7603     7682      +79     
==========================================
- Hits         6429     5828     -601     
- Misses       1174     1854     +680     

Flag	Coverage Δ
unittests	75.86% <91.66%> (-8.70%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files	Coverage Δ
sbi/inference/posteriors/mcmc_posterior.py	85.93% <91.66%> (-0.26%)	⬇️

... and 28 files with indirect coverage changes

[gmoss13]

Thanks a lot @janfb, I think these changes fix the issues we discussed! I added some suggestions, happy to discuss further (also happy to implement them myself if you agree).

[gmoss13]

Thanks Jan! Added a couple of very minor comments.