Repository: pymc-devs/pymc
Pull request: Fix error in `dataset_to_point_list` when chain, draw are not the leading dims
## Description
Split the dict comprehension in two to keep transposing and reshaping coherent.
Otherwise, the sizes used in the reshaping are taken from the object before transposing
and the function doesn't work whenever the transposing is actually needed.

## Related Issue
<!--- It is good practice to first open an issue explaining the bug / new feature that is addressed by this PR -->
<!--- Please type an `x` in one of the boxes below and provide the issue number after the # sign: -->
- [x] Closes #7178

## Checklist
<!--- Make sure you have completed the following steps before submitting your PR -->
<!--- Feel free to type an `x` in all the boxes below to let us know you have completed the steps: -->
- [x] Checked that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html)
- [x] Included tests that prove the fix is effective or that the new feature works
- [x] If you are a pro: each commit corresponds to a [relevant logical change](https://wiki.openstack.org/wiki/GitCommitMessages#Structural_split_of_changes)
<!--- You may find this guide helpful: https://mainmatter.com/blog/2021/05/26/keeping-a-clean-git-history/ -->

## Type of change
<!--- Select one of the categories below by typing an `x` in the box -->
- [ ] New feature / enhancement
- [x] Bug fix
- [ ] Documentation
- [ ] Maintenance
- [ ] Other (please specify):
<!--- Additionally, if you are a maintainer or reviewer, please make sure that the appropriate labels are added to this PR -->


<!-- readthedocs-preview pymc start -->
----
📚 Documentation preview 📚: https://pymc--7180.org.readthedocs.build/en/7180/

<!-- readthedocs-preview pymc end -->

Problem statement:
`dataset_to_point_list` fails when chain, draw are not the leading dims
### Description

```python
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

Hints from issue discussion:
I am actually surprised it breaks there and not in the reshaping part. The culprit is clearly this line:

https://github.com/pymc-devs/pymc/blob/6252d2e58dc211c913ee2e652a4058d271d48bbd/pymc/util.py#L253

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.

Commit message: fix dataset_to_point_list

Patch to review:
diff --git a/pymc/util.py b/pymc/util.py
index 799d92b1b4..1062f3f72f 100644
--- a/pymc/util.py
+++ b/pymc/util.py
@@ -249,9 +249,9 @@ def dataset_to_point_list(
             raise ValueError(f"Variable names must be str, but dataset key {vn} is a {type(vn)}.")
     num_sample_dims = len(sample_dims)
     stacked_dims = {dim_name: ds[var_names[0]][dim_name] for dim_name in sample_dims}
+    stacked_dict = {vn: da.transpose(*sample_dims, ...) for vn, da in ds.items()}
     stacked_dict = {
-        vn: da.transpose(*sample_dims, ...).values.reshape((-1, *da.shape[num_sample_dims:]))
-        for vn, da in ds.items()
+        vn: da.values.reshape((-1, *da.shape[num_sample_dims:])) for vn, da in ds.items()
     }
     points = [
         {vn: stacked_dict[vn][i, ...] for vn in var_names}
diff --git a/tests/test_util.py b/tests/test_util.py
index e984fdbd1c..61d916249e 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -170,6 +170,16 @@ def test_dataset_to_point_list(input_type):
     assert isinstance(pl[0]["A"], np.ndarray)
 
 
+def test_transposed_dataset_to_point_list():
+    ds = xarray.Dataset()
+    ds["A"] = xarray.DataArray([[[1, 2, 3], [2, 3, 4]]] * 5, dims=("team", "draw", "chain"))
+    pl, _ = dataset_to_point_list(ds, sample_dims=["chain", "draw"])
+    assert isinstance(pl, list)
+    assert len(pl) == 6
+    assert isinstance(pl[0], dict)
+    assert isinstance(pl[0]["A"], np.ndarray)
+
+
 def test_dataset_to_point_list_str_key():
     # Check that non-str keys are caught
     ds = xarray.Dataset()


Code review comments:
diff_hunk: @@ -249,9 +249,9 @@ def dataset_to_point_list(
             raise ValueError(f"Variable names must be str, but dataset key {vn} is a {type(vn)}.")
     num_sample_dims = len(sample_dims)
     stacked_dims = {dim_name: ds[var_names[0]][dim_name] for dim_name in sample_dims}
+    stacked_dict = {vn: da.transpose(*sample_dims, ...) for vn, da in ds.items()}
     stacked_dict = {
-        vn: da.transpose(*sample_dims, ...).values.reshape((-1, *da.shape[num_sample_dims:]))
-        for vn, da in ds.items()
+        vn: da.values.reshape((-1, *da.shape[num_sample_dims:])) for vn, da in ds.items()
path: pymc/util.py
text: @author:
```suggestion
        vn: da.values.reshape((-1, *da.shape[num_sample_dims:])) for vn, da in transposed_dict.items()
```

diff_hunk: @@ -249,9 +249,9 @@ def dataset_to_point_list(
             raise ValueError(f"Variable names must be str, but dataset key {vn} is a {type(vn)}.")
     num_sample_dims = len(sample_dims)
     stacked_dims = {dim_name: ds[var_names[0]][dim_name] for dim_name in sample_dims}
+    stacked_dict = {vn: da.transpose(*sample_dims, ...) for vn, da in ds.items()}
path: pymc/util.py
text: @author:
```suggestion
    transposed_dict = {vn: da.transpose(*sample_dims, ...) for vn, da in ds.items()}
```