Add typing in assign_step_methods

Author: Joseph Hall

pymc/sampling/mcmc.py

@@ -20,15 +20,16 @@
 import time
 import warnings
 
-from collections import defaultdict
 from typing import (
     Any,
     Dict,
     Iterator,
     List,
     Literal,
+    Mapping,
     Optional,
     Sequence,
+    Set,
     Tuple,
     Type,
     Union,
@@ -40,6 +41,7 @@
 
 from arviz import InferenceData
 from fastprogress.fastprogress import progress_bar
+from pytensor.graph.basic import Variable
 from typing_extensions import Protocol, TypeAlias
 
 import pymc as pm
@@ -91,9 +93,9 @@ def __call__(self, trace: IBaseTrace, draw: Draw):
 
 
 def instantiate_steppers(
-    model,
+    model: Model,
     steps: List[Step],
-    selected_steps: Dict[Type[BlockedStep], List[Any]],
+    selected_steps: Mapping[Type[BlockedStep], List[Any]],
     step_kwargs: Optional[Dict[str, Dict]] = None,
 ) -> Union[Step, List[Step]]:
     """Instantiate steppers assigned to the model variables.
@@ -149,7 +151,12 @@ def instantiate_steppers(
     return steps
 
 
-def assign_step_methods(model, step=None, methods=None, step_kwargs=None):
+def assign_step_methods(
+    model: Model,
+    step: Optional[Union[Step, Sequence[Step]]] = None,
+    methods: Optional[Sequence[Type[BlockedStep]]] = None,
+    step_kwargs: Optional[Dict[str, Any]] = None,
+) -> Union[Step, List[Step]]:
     """Assign model variables to appropriate step methods.
 
     Passing a specified model will auto-assign its constituent stochastic
@@ -179,49 +186,48 @@ def assign_step_methods(model, step=None, methods=None, step_kwargs=None):
     methods : list
         List of step methods associated with the model's variables.
     """
-    steps = []
-    assigned_vars = set()
-
-    if methods is None:
-        methods = pm.STEP_METHODS
+    steps: List[Step] = []
+    assigned_vars: Set[Variable] = set()
 
     if step is not None:
-        try:
-            steps += list(step)
-        except TypeError:
+        if isinstance(step, (BlockedStep, CompoundStep)):
             steps.append(step)
+        else:
+            steps.extend(step)
         for step in steps:
             for var in step.vars:
                 if var not in model.value_vars:
                     raise ValueError(
-                        f"{var} assigned to {step} sampler is not a value variable in the model. You can use `util.get_value_vars_from_user_vars` to parse user provided variables."
+                        f"{var} assigned to {step} sampler is not a value variable in the model. "
+                        "You can use `util.get_value_vars_from_user_vars` to parse user provided variables."
                     )
             assigned_vars = assigned_vars.union(set(step.vars))
 
     # Use competence classmethods to select step methods for remaining
     # variables
-    selected_steps = defaultdict(list)
+    methods_list: List[Type[BlockedStep]] = list(methods or pm.STEP_METHODS)
+    selected_steps: Dict[Type[BlockedStep], List] = {}
     model_logp = model.logp()
 
     for var in model.value_vars:
         if var not in assigned_vars:
             # determine if a gradient can be computed
-            has_gradient = var.dtype not in discrete_types
+            has_gradient = getattr(var, "dtype") not in discrete_types
             if has_gradient:
                 try:
-                    tg.grad(model_logp, var)
+                    tg.grad(model_logp, var)  # type: ignore
                 except (NotImplementedError, tg.NullTypeGradError):
                     has_gradient = False
 
             # select the best method
             rv_var = model.values_to_rvs[var]
             selected = max(
-                methods,
-                key=lambda method, var=rv_var, has_gradient=has_gradient: method._competence(
+                methods_list,
+                key=lambda method, var=rv_var, has_gradient=has_gradient: method._competence(  # type: ignore
                     var, has_gradient
                 ),
             )
-            selected_steps[selected].append(var)
+            selected_steps.setdefault(selected, []).append(var)
 
     return instantiate_steppers(model, steps, selected_steps, step_kwargs)
 

pymc/step_methods/compound.py

@@ -246,7 +246,7 @@ def reset_tuning(self):
                 method.reset_tuning()
 
     @property
-    def vars(self):
+    def vars(self) -> List[Variable]:
         return [var for method in self.methods for var in method.vars]