Fix several typing gaps

Author: michaelosthege

pymc/sampling/mcmc.py

@@ -21,14 +21,14 @@
 import warnings
 
 from collections import defaultdict
-from typing import Iterator, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import pytensor.gradient as tg
 
 from arviz import InferenceData
 from fastprogress.fastprogress import progress_bar
-from typing_extensions import TypeAlias
+from typing_extensions import Protocol, TypeAlias
 
 import pymc as pm
 
@@ -64,6 +64,13 @@
 Step: TypeAlias = Union[BlockedStep, CompoundStep]
 
 
+class SamplingIteratorCallback(Protocol):
+    """Signature of the callable that may be passed to `pm.sample(callable=...)`."""
+
+    def __call__(self, trace: BaseTrace, draw: Draw):
+        pass
+
+
 _log = logging.getLogger("pymc")
 
 
@@ -221,7 +228,7 @@ def sample(
     cores: Optional[int] = None,
     tune: int = 1000,
     progressbar: bool = True,
-    model=None,
+    model: Optional[Model] = None,
     random_seed: RandomState = None,
     discard_tuned_samples: bool = True,
     compute_convergence_checks: bool = True,
@@ -599,7 +606,7 @@ def sample(
 
     idata = None
     if compute_convergence_checks or return_inferencedata:
-        ikwargs = dict(model=model, save_warmup=not discard_tuned_samples)
+        ikwargs: Dict[str, Any] = dict(model=model, save_warmup=not discard_tuned_samples)
         if idata_kwargs:
             ikwargs.update(idata_kwargs)
         idata = pm.to_inference_data(mtrace, **ikwargs)
@@ -655,8 +662,8 @@ def _sample_many(
     traces: Sequence[BaseTrace],
     start: Sequence[PointType],
     random_seed: Optional[Sequence[RandomSeed]],
-    step,
-    callback=None,
+    step: Step,
+    callback: Optional[SamplingIteratorCallback] = None,
     **kwargs,
 ):
     """Samples all chains sequentially.
@@ -695,7 +702,7 @@ def _sample(
     random_seed: RandomSeed,
     start: PointType,
     draws: int,
-    step=None,
+    step: Step,
     trace: BaseTrace,
     tune: int,
     model: Optional[Model] = None,
@@ -760,14 +767,14 @@ def _sample(
 def _iter_sample(
     *,
     draws: int,
-    step,
+    step: Step,
     start: PointType,
     trace: BaseTrace,
     chain: int = 0,
     tune: int = 0,
-    model=None,
+    model: Optional[Model] = None,
     random_seed: RandomSeed = None,
-    callback=None,
+    callback: Optional[SamplingIteratorCallback] = None,
 ) -> Iterator[bool]:
     """Generator for sampling one chain. (Used in singleprocess sampling.)
 
@@ -803,19 +810,13 @@ def _iter_sample(
     if random_seed is not None:
         np.random.seed(random_seed)
 
-    try:
-        step = CompoundStep(step)
-    except TypeError:
-        pass
-
     point = start
 
     try:
         step.tune = bool(tune)
         if hasattr(step, "reset_tuning"):
             step.reset_tuning()
         for i in range(draws):
-            stats = None
             diverging = False
 
             if i == 0 and hasattr(step, "iter_count"):
@@ -825,7 +826,7 @@ def _iter_sample(
             point, stats = step.step(point)
             trace.record(point, stats)
             log_warning_stats(stats)
-            diverging = i > tune and stats and stats[0].get("diverging")
+            diverging = i > tune and len(stats) > 0 and (stats[0].get("diverging") == True)
             if callback is not None:
                 callback(
                     trace=trace,
@@ -854,8 +855,8 @@ def _mp_sample(
     start: Sequence[PointType],
     progressbar: bool = True,
     traces: Sequence[BaseTrace],
-    model=None,
-    callback=None,
+    model: Optional[Model] = None,
+    callback: Optional[SamplingIteratorCallback] = None,
     mp_ctx=None,
     **kwargs,
 ) -> None:
@@ -884,7 +885,7 @@ def _mp_sample(
         A backend instance, or None.
         If None, the NDArray backend is used.
     model : Model (optional if in ``with`` context)
-    callback : Callable
+    callback
         A function which gets called for every sample from the trace of a chain. The function is
         called with the trace and the current draw and will contain all samples for a single trace.
         the ``draw.chain`` argument can be used to determine which of the active chains the sample
@@ -994,7 +995,7 @@ def init_nuts(
     init: str = "auto",
     chains: int = 1,
     n_init: int = 500_000,
-    model=None,
+    model: Optional[Model] = None,
     random_seed: RandomSeed = None,
     progressbar=True,
     jitter_max_retries: int = 10,

pymc/sampling/population.py

@@ -28,7 +28,7 @@
 
 from pymc.backends.base import BaseTrace
 from pymc.initial_point import PointType
-from pymc.model import modelcontext
+from pymc.model import Model, modelcontext
 from pymc.stats.convergence import log_warning_stats
 from pymc.step_methods import CompoundStep
 from pymc.step_methods.arraystep import (
@@ -50,13 +50,13 @@
 
 def _sample_population(
     *,
-    initial_points,
+    initial_points: Sequence[PointType],
     draws: int,
     start: Sequence[PointType],
     random_seed: RandomSeed,
-    step,
+    step: Union[BlockedStep, CompoundStep],
     tune: int,
-    model,
+    model: Model,
     progressbar: bool = True,
     parallelize: bool = False,
     traces: Sequence[BaseTrace],
@@ -108,7 +108,14 @@ def _sample_population(
     return
 
 
-def warn_population_size(*, step: CompoundStep, initial_points, model, chains: int):
+def warn_population_size(
+    *,
+    step: Union[BlockedStep, CompoundStep],
+    initial_points: Sequence[PointType],
+    model: Model,
+    chains: int,
+):
+    """Emit informative errors/warnings for dangerously small population size."""
     has_demcmc = np.any(
         [
             isinstance(m, DEMetropolis)

pymc/stats/log_likelihood.py

@@ -11,7 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-from typing import Optional, Sequence
+from typing import Optional, Sequence, cast
 
 import numpy as np
 
@@ -22,6 +22,7 @@
 
 from pymc.backends.arviz import _DefaultTrace
 from pymc.model import Model, modelcontext
+from pymc.pytensorf import PointFunc
 from pymc.util import dataset_to_point_list
 
 __all__ = ("compute_log_likelihood",)
@@ -86,6 +87,7 @@ def compute_log_likelihood(
             outs=model.logp(vars=observed_vars, sum=False),
             on_unused_input="ignore",
         )
+        elemwise_loglike_fn = cast(PointFunc, elemwise_loglike_fn)
     finally:
         model.rvs_to_values = original_rvs_to_values
         model.rvs_to_transforms = original_rvs_to_transforms

pymc/step_methods/compound.py

@@ -36,6 +36,7 @@ def __init__(self, methods):
         self.name = (
             f"Compound[{', '.join(getattr(m, 'name', 'UNNAMED_STEP') for m in self.methods)}]"
         )
+        self.tune = True
 
     def step(self, point) -> Tuple[PointType, StatsType]:
         stats = []

pymc/util.py

@@ -235,7 +235,7 @@ def enhanced(*args, **kwargs):
 
 
 def dataset_to_point_list(
-    ds: xarray.Dataset, sample_dims: List
+    ds: xarray.Dataset, sample_dims: Sequence[str]
 ) -> Tuple[List[Dict[str, np.ndarray]], Dict[str, Any]]:
     # All keys of the dataset must be a str
     var_names = list(ds.keys())

scripts/run_mypy.py

@@ -42,7 +42,6 @@
 pymc/printing.py
 pymc/pytensorf.py
 pymc/sampling/jax.py
-pymc/stats/log_likelihood.py
 pymc/variational/approximations.py
 pymc/variational/opvi.py
 """