Bump aesara to 2.7.8. (#5995)

* Bump aesara to 2.7.8.

* Bump pre-commit aesara.

* Ignore mypy typing errors for pymc/util.py

* Add safety check for `xarray.Dataset` key types

Co-authored-by: Michael Osthege <[email protected]>

Author: twiecki

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
       - types-filelock
       - types-setuptools
       - arviz
-      - aesara==2.7.7
+      - aesara==2.7.8
       - aeppl==0.0.32
       always_run: true
       require_serial: true

conda-envs/environment-dev.yml

@@ -6,7 +6,7 @@ channels:
 dependencies:
 # Base dependencies
 - aeppl=0.0.32
-- aesara=2.7.7
+- aesara=2.7.8
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

conda-envs/environment-test.yml

@@ -6,7 +6,7 @@ channels:
 dependencies:
 # Base dependencies
 - aeppl=0.0.32
-- aesara=2.7.7
+- aesara=2.7.8
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

conda-envs/windows-environment-dev.yml

@@ -6,7 +6,7 @@ channels:
 dependencies:
 # Base dependencies (see install guide for Windows)
 - aeppl=0.0.32
-- aesara=2.7.7
+- aesara=2.7.8
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

conda-envs/windows-environment-test.yml

@@ -6,7 +6,7 @@ channels:
 dependencies:
 # Base dependencies (see install guide for Windows)
 - aeppl=0.0.32
-- aesara=2.7.7
+- aesara=2.7.8
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

pymc/tests/test_util.py

@@ -14,13 +14,20 @@
 
 import numpy as np
 import pytest
+import xarray
 
 from cachetools import cached
 
 import pymc as pm
 
 from pymc.distributions.transforms import RVTransform
-from pymc.util import UNSET, hash_key, hashable, locally_cachedmethod
+from pymc.util import (
+    UNSET,
+    dataset_to_point_list,
+    hash_key,
+    hashable,
+    locally_cachedmethod,
+)
 
 
 class TestTransformName:
@@ -142,3 +149,18 @@ def fn(a=UNSET):
     help(fn)
     captured = capsys.readouterr()
     assert "a=UNSET" in captured.out
+
+
+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)
+    assert isinstance(pl, list)
+    assert len(pl) == 6
+    assert isinstance(pl[0], dict)
+    assert isinstance(pl[0]["A"], np.ndarray)
+
+    # 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)

pymc/util.py

@@ -14,7 +14,7 @@
 
 import functools
 
-from typing import Dict, List, Tuple, Union, cast
+from typing import Dict, Hashable, List, Tuple, Union, cast
 
 import arviz
 import cloudpickle
@@ -232,15 +232,18 @@ def enhanced(*args, **kwargs):
 
 
 def dataset_to_point_list(ds: xarray.Dataset) -> List[Dict[str, np.ndarray]]:
+    # All keys of the dataset must be a str
+    for vn in ds.keys():
+        if not isinstance(vn, str):
+            raise ValueError(f"Variable names must be str, but dataset key {vn} is a {type(vn)}.")
     # make dicts
-    points: List[Dict[str, np.ndarray]] = []
-    vn: str
+    points: List[Dict[Hashable, np.ndarray]] = []
     da: "xarray.DataArray"
     for c in ds.chain:
         for d in ds.draw:
             points.append({vn: da.sel(chain=c, draw=d).values for vn, da in ds.items()})
     # use the list of points
-    return points
+    return cast(List[Dict[str, np.ndarray]], points)
 
 
 def chains_and_samples(data: Union[xarray.Dataset, arviz.InferenceData]) -> Tuple[int, int]:

requirements-dev.txt

@@ -2,7 +2,7 @@
 # See that file for comments about the need/usage of each dependency.
 
 aeppl==0.0.32
-aesara==2.7.7
+aesara==2.7.8
 arviz>=0.12.0
 cachetools>=4.2.1
 cloudpickle

requirements.txt

@@ -1,5 +1,5 @@
 aeppl==0.0.32
-aesara==2.7.7
+aesara==2.7.8
 arviz>=0.12.0
 cachetools>=4.2.1
 cloudpickle