Make trace backend related type hints more generic

* Specify covariant input types in `StatsBijection`.
* Annotate `_choose_chains` to be independent of `BaseTrace` type.

Author: michaelosthege

pymc/backends/base.py

@@ -21,7 +21,18 @@
 import warnings
 
 from abc import ABC
-from typing import Dict, List, Optional, Sequence, Set, Tuple, Union, cast
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Sized,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import numpy as np
 
@@ -510,7 +521,10 @@ def _squeeze_cat(results, combine, squeeze):
     return results
 
 
-def _choose_chains(traces: Sequence[BaseTrace], tune: int) -> Tuple[List[BaseTrace], int]:
+S = TypeVar("S", bound=Sized)
+
+
+def _choose_chains(traces: Sequence[S], tune: int) -> Tuple[List[S], int]:
     """
     Filter and slice traces such that (n_traces * len(shortest_trace)) is maximized.
 

pymc/step_methods/compound.py

@@ -20,7 +20,7 @@
 
 from abc import ABC, abstractmethod
 from enum import IntEnum, unique
-from typing import Dict, List, Sequence, Tuple, Union
+from typing import Any, Dict, List, Mapping, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -181,14 +181,14 @@ def flatten_steps(step: Union[BlockedStep, CompoundStep]) -> List[BlockedStep]:
 class StatsBijection:
     """Map between a `list` of stats to `dict` of stats."""
 
-    def __init__(self, sampler_stats_dtypes: Sequence[Dict[str, type]]) -> None:
+    def __init__(self, sampler_stats_dtypes: Sequence[Mapping[str, type]]) -> None:
         # Keep a list of flat vs. original stat names
         self._stat_groups: List[List[Tuple[str, str]]] = [
             [(f"sampler_{s}__{statname}", statname) for statname, _ in names_dtypes.items()]
             for s, names_dtypes in enumerate(sampler_stats_dtypes)
         ]
 
-    def map(self, stats_list: StatsType) -> StatsDict:
+    def map(self, stats_list: Sequence[Mapping[str, Any]]) -> StatsDict:
         """Combine stats dicts of multiple samplers into one dict."""
         stats_dict = {}
         for s, sts in enumerate(stats_list):
@@ -197,7 +197,7 @@ def map(self, stats_list: StatsType) -> StatsDict:
                 stats_dict[sname] = sval
         return stats_dict
 
-    def rmap(self, stats_dict: StatsDict) -> StatsType:
+    def rmap(self, stats_dict: Mapping[str, Any]) -> StatsType:
         """Split a global stats dict into a list of sampler-wise stats dicts."""
         stats_list = []
         for namemap in self._stat_groups: