Add option to keep chains unflattened.

[tillahoffmann]

This PR

• adds a function get_samples(param: str, flatten_chains: bool = True) -> np.ndarray to the Fit object which reproduces the behaviour of __getitem__ by default. However, setting flatten_chains=False leaves the trailing dimension as (num_saved_samples, num_chains). This behaviour is useful for comparing initialisations of different chains, computing Rhat, etc.
• adds a chain__ column to the data frame produced by Fit.to_frame which can be used to identify which chain generated a given sample.

[riddell-stan]

Thanks. Looks interesting.

We're only using GitHub right now for bugs and confirmed features (i.e., features which have been discussed on the Stan forums). So the thing to do would be to post there first.

[riddell-stan]

I'm going to close this for now. Adding features that are not currently in RStan (including the hypothetical RStan 3) is out of scope right now.

I also think you might be able to implement this with a plugin.