util.py related fixes

Author: OriolAbril

pymc/tests/test_util.py

@@ -154,7 +154,7 @@ def fn(a=UNSET):
 def test_dataset_to_point_list():
     ds = xarray.Dataset()
     ds["A"] = xarray.DataArray([[1, 2, 3]] * 2, dims=("chain", "draw"))
-    pl = dataset_to_point_list(ds)
+    pl, _ = dataset_to_point_list(ds, sample_dims=["chain", "draw"])
     assert isinstance(pl, list)
     assert len(pl) == 6
     assert isinstance(pl[0], dict)
@@ -163,4 +163,4 @@ def test_dataset_to_point_list():
     # Check that non-str keys are caught
     ds[3] = xarray.DataArray([1, 2, 3])
     with pytest.raises(ValueError, match="must be str"):
-        dataset_to_point_list(ds)
+        dataset_to_point_list(ds, sample_dims=["chain", "draw"])

pymc/util.py

@@ -14,7 +14,7 @@
 
 import functools
 
-from typing import Dict, List, cast
+from typing import Any, Dict, List, Tuple, cast
 
 import cloudpickle
 import numpy as np
@@ -230,7 +230,9 @@ def enhanced(*args, **kwargs):
     return enhanced
 
 
-def dataset_to_point_list(ds: xarray.Dataset, sample_dims: List) -> List[Dict[str, np.ndarray]]:
+def dataset_to_point_list(
+    ds: xarray.Dataset, sample_dims: List
+) -> Tuple[List[Dict[str, np.ndarray]], Dict[str, Any]]:
     # All keys of the dataset must be a str
     var_names = list(ds.keys())
     for vn in var_names: