`.from_dict()` messes up chain/draw dimensions when chains have uneven lengths

[michaelosthege]

Describe the bug
When the chains fed to .from_dict() have different lengths, the chain and draw dimension collapse to chain: 1, draw: nchains.

To Reproduce

import arviz
import numpy as np
array = np.array

kwargs = {
    'warmup_posterior': {
        'A': [
            array([0, 1, 2, 3, 4], dtype=np.uint16),
            array([0, 1, 2, 3, 4], dtype=np.uint16)
        ]
    },
    'warmup_sample_stats': {
        'tune': [
            array([ True,  True,  True,  True,  True]),
            array([ True,  True,  True,  True,  True])
        ]
    },
    'posterior': {
        'A': [
            array([ 5,  6,  7,  8,  9, 10, 11], dtype=np.uint16),
            array([ 5,  6,  7,  8,  9, 10], dtype=np.uint16)
        ]
    },
    'sample_stats': {
        'tune': [
            array([False, False, False, False, False, False, False]),
            array([False, False, False, False, False, False])
        ]
    },
    'coords': {},
    'dims': {
        'A': []
    },
    'attrs': {},
    'constant_data': {},
    'observed_data': {},
    'save_warmup': True
}
idata = arviz.from_dict(**kwargs)
print(idata.warmup_posterior.dims) # {'chain': 2, 'draw': 5}
print(idata.posterior.dims)        # {'chain': 1, 'draw': 2} but shoud be {'chain': 2, 'draw': 7}

Expected behavior
I expected the "missing" entries of the shorter chain to be filled-up by np.nan such that both chains have the same length (equal to the length of the longer chain).

Additional context
Looking at the idata.posterior, the "A" variable is stored as

array(
    [
        [array([ 5,  6,  7,  8,  9, 10, 11], dtype=uint16),
         array([ 5,  6,  7,  8,  9, 10], dtype=uint16)]
    ],
    dtype=object
)

which is an object-array containing one "chain" that has two "draws" each of which is the vector of draws of the actual chains.

This bug was the underlying cause of pymc-devs/mcbackend#44

[Chirag3841]

@michaelosthege I guess NumPy converts ragged nested arrays into dtype=object arrays, and arviz.from_dict() does not detect this situation.
That' s why chain inference is incorrect.
We can improve that by Pad shorter chains to the longest chain length.

[michaelosthege]

Yes, but that's not desirable.

I ran into this problem again in pymc-devs/mcbackend#128 - it's getting worse: With ArviZ 1.0 changes my workaround that I had implemented in McBackend no longer works.

[Chirag3841]

@michaelosthege
Other than that, best fix is that
before calling arviz.from_dict(), detect ragged chains and raise an error like ragged chains detected which means chains length differ.

[michaelosthege]

Different chain lenghts should not be considered errornous. There is no requirement from the math that independent MCMC chains must have equal lengths. It is purely a limitation of how the arrays are currently stored in ArviZ/xarray.

[Chirag3841]

@michaelosthege Yes you are right
different chain lengths are not invalid. There is no mathematical requirement that independent MCMC chains must have equal length.
This is purely a limitation of how ArviZ stores arrays.
Therefore ArviZ should either pad shorter chains or explicitly raise an error, but it must not silently collapse the chain and draw dimensions since that corrupts the posterior structure.