dataset_to_point_list fails when chain, draw are not the leading dims

[ricardoV94]

Description

import pymc as pm

with pm.Model(coords={"trial": [0]}) as m:
    x = pm.Normal("x", shape=(1,), dims="trial")
    y = pm.Normal("y", x, observed=[5], dims="trial")
    
    idata = pm.sample(tune=0, draws=10)
    idata = idata.posterior.transpose("chain", "trial", "draw")
    pm.sample_posterior_predictive(idata)  # IndexError: index 2 is out of bounds for axis 0 with size 2

CC @OriolAbril

[OriolAbril]

I am actually surprised it breaks there and not in the reshaping part. The culprit is clearly this line:

pymc/pymc/util.py

Line 253 in 6252d2e

vn: da.transpose(*sample_dims, ...).values.reshape((-1, *da.shape[num_sample_dims:]))

we transpose the dataarray to have everything work independently of the dimension order, but then the reshape uses the yet-to-be-transposed dataarray to get the shape of the final array.