Integration with InferenceObjects.jl

[sethaxen]

On Twitter, @yebai suggested adding integration with InferenceObjects to MCMCChains: https://twitter.com/Hong_Ge2/status/1560343482216103938. I'm opening this issue for further discussion.

InferenceObjects.InferenceData is the storage format for Monte Carlo draws used by ArviZ.jl. Along with Python's arviz.InferenceData, it follows the cross-language InferenceData schema. PyMC uses Python's implementation as its official sample storage format. InferenceData can be serialized to NetCDF to standardize communicating results of Bayesian analyses across languages and PPLs. In Julia, it is built on DimensionalData. See example usage and plotting examples (using the Tables interface).

@yebai's suggestion is ultimately to deprecate Chains to instead use InferenceData. I see several upsides of this approach:

1. Chains is based on the somewhat outdated AxisArrays, while DimensionalData is more modern.
2. Chains flattens all draws and sampling statistics into a single 3D float array, which discards a lot of the structure of the sampled types (which may themselves be multidimensional or have non-float eltypes, such as Int or even Cholesky).
3. InferenceData's features are a superset of Chains. It can get closer to the original structure of the user's samples with named dimensions, but it also supports storing other metadata and can store prior, predictive, log-likelihood, and warmup draws, as well as the original data.
4. InferenceObjects is a relatively light dependency (~0.120-0.2s load time on Julia v1.7-1.8 vs MCMCChains with 1.7-3.6s) so would not add much to MCMCChains's load time.

Currently ArviZ.jl has a converter from_mcmcchains, which is used to convert Chains to InferenceData. Integration between Chains and InferenceData might look like the following steps:

1. Move ArviZ.from_mcmcchains here (with a better name)
2. Make InferenceData a supported chain_type for AbstractMCMC.sample (https://beta.turing.ml/AbstractMCMC.jl/dev/api/#Chains), which would bypass Chains's flattening entirely. I'm not sure this should live here, but it should not live in InferenceObjects.

[sethaxen]

2. Make InferenceData a supported chain_type for AbstractMCMC.sample (https://beta.turing.ml/AbstractMCMC.jl/dev/api/#Chains), which would bypass Chains's flattening entirely. I'm not sure this should live here, but it should not live in InferenceObjects.

It looks like this would involved implementing AbstractMCMC.bundle_samples for specific samplers. For Chains this is done in Turing itself, not here, e.g. https://github.com/TuringLang/Turing.jl/blob/9f8a9c4c476095d45246b4924d5cd542f3f8d506/src/inference/Inference.jl#L320-L391

[sethaxen]

Since there have been no objections to these steps, I'm going to move forward with opening a PR for Step 1.

[cpfiffer]

Okay, thank you!

[sethaxen]

I wonder actually if this is going the wrong way about this. MCMCChains destructively flattens draws into one large array, and a converter then needs to infer from the variable names how to unflatten the draws, whereas Turing can return NamedTuples containing unflattened draws. It might be better to just directly implement DynamicPPL/AbstractMCMC interfaces for InferenceData storage, so instead of converting a Chains to an InferenceData, anyone sampling with those interfaces can just build the InferenceData directly.

I've opened a draft PR at TuringLang/Turing.jl#1913