Add type hints to SMC code

Author: michaelosthege

pymc/smc/kernels.py

@@ -16,14 +16,15 @@
 import warnings
 
 from abc import ABC
-from typing import Dict, cast
+from typing import Dict, Union, cast
 
 import numpy as np
 import pytensor.tensor as at
 
 from pytensor.graph.replace import clone_replace
 from scipy.special import logsumexp
 from scipy.stats import multivariate_normal
+from typing_extensions import TypeAlias
 
 from pymc.backends.ndarray import NDArray
 from pymc.blocking import DictToArrayBijection
@@ -39,6 +40,9 @@
 from pymc.step_methods.metropolis import MultivariateNormalProposal
 from pymc.vartypes import discrete_types
 
+SMCStats: TypeAlias = Dict[str, Union[int, float]]
+SMCSettings: TypeAlias = Dict[str, Union[int, float]]
+
 
 class SMC_KERNEL(ABC):
     """Base class for the Sequential Monte Carlo kernels.
@@ -304,7 +308,7 @@ def mutate(self):
         """Apply kernel-specific perturbation to the particles once per stage"""
         pass
 
-    def sample_stats(self) -> Dict:
+    def sample_stats(self) -> SMCStats:
         """Stats to be saved at the end of each stage
 
         These stats will be saved under `sample_stats` in the final InferenceData object.
@@ -314,7 +318,7 @@ def sample_stats(self) -> Dict:
             "beta": self.beta,
         }
 
-    def sample_settings(self) -> Dict:
+    def sample_settings(self) -> SMCSettings:
         """SMC_kernel settings to be saved once at the end of sampling.
 
         These stats will be saved under `sample_stats` in the final InferenceData object.
@@ -425,7 +429,7 @@ def mutate(self):
 
         self.acc_rate = np.mean(ac_)
 
-    def sample_stats(self):
+    def sample_stats(self) -> SMCStats:
         stats = super().sample_stats()
         stats.update(
             {
@@ -434,7 +438,7 @@ def sample_stats(self):
         )
         return stats
 
-    def sample_settings(self):
+    def sample_settings(self) -> SMCSettings:
         stats = super().sample_settings()
         stats.update(
             {
@@ -543,7 +547,7 @@ def mutate(self):
 
         self.chain_acc_rate = np.mean(ac_, axis=0)
 
-    def sample_stats(self):
+    def sample_stats(self) -> SMCStats:
         stats = super().sample_stats()
         stats.update(
             {
@@ -553,7 +557,7 @@ def sample_stats(self):
         )
         return stats
 
-    def sample_settings(self):
+    def sample_settings(self) -> SMCSettings:
         stats = super().sample_settings()
         stats.update(
             {

pymc/smc/sampling.py

@@ -19,6 +19,7 @@
 
 from collections import defaultdict
 from itertools import repeat
+from typing import Any, Dict, Optional, Tuple, Union
 
 import cloudpickle
 import numpy as np
@@ -30,7 +31,7 @@
 
 from pymc.backends.arviz import dict_to_dataset, to_inference_data
 from pymc.backends.base import MultiTrace
-from pymc.model import modelcontext
+from pymc.model import Model, modelcontext
 from pymc.sampling.parallel import _cpu_count
 from pymc.smc.kernels import IMH
 from pymc.util import RandomState, _get_seeds_per_chain
@@ -50,7 +51,7 @@ def sample_smc(
     idata_kwargs=None,
     progressbar=True,
     **kernel_kwargs,
-):
+) -> Union[InferenceData, MultiTrace]:
     r"""
     Sequential Monte Carlo based sampling.
 
@@ -237,19 +238,23 @@ def sample_smc(
 
     if compute_convergence_checks:
         _compute_convergence_checks(idata, draws, model, trace)
-    return idata if return_inferencedata else trace
+
+    if return_inferencedata:
+        assert idata is not None
+        return idata
+    return trace
 
 
 def _save_sample_stats(
     sample_settings,
     sample_stats,
     chains,
-    trace,
-    return_inferencedata,
+    trace: MultiTrace,
+    return_inferencedata: bool,
     _t_sampling,
     idata_kwargs,
-    model,
-):
+    model: Model,
+) -> Tuple[Optional[Any], Optional[InferenceData]]:
     sample_settings_dict = sample_settings[0]
     sample_settings_dict["_t_sampling"] = _t_sampling
     sample_stats_dict = sample_stats[0]
@@ -262,12 +267,12 @@ def _save_sample_stats(
                 value_list.append(chain_sample_stats[stat])
             sample_stats_dict[stat] = value_list
 
+    idata: Optional[InferenceData] = None
     if not return_inferencedata:
         for stat, value in sample_stats_dict.items():
             setattr(trace.report, stat, value)
         for stat, value in sample_settings_dict.items():
             setattr(trace.report, stat, value)
-        idata = None
     else:
         for stat, value in sample_stats_dict.items():
             if chains > 1:
@@ -284,7 +289,7 @@ def _save_sample_stats(
             library=pymc,
         )
 
-        ikwargs = dict(model=model)
+        ikwargs: Dict[str, Any] = dict(model=model)
         if idata_kwargs is not None:
             ikwargs.update(idata_kwargs)
         idata = to_inference_data(trace, **ikwargs)
@@ -293,7 +298,9 @@ def _save_sample_stats(
     return sample_stats, idata
 
 
-def _compute_convergence_checks(idata, draws, model, trace):
+def _compute_convergence_checks(
+    idata: Optional[InferenceData], draws: int, model: Model, trace: MultiTrace
+):
     if draws < 100:
         warnings.warn(
             "The number of samples is too small to check convergence reliably.",