Update type annotations for Pyro 1.9

.github/workflows/lint.yml

@@ -12,6 +12,7 @@ jobs:
     strategy:
       matrix:
         python-version: [3.8, 3.9, '3.10']
+        pyro-version: ['1.8.6', 'latest']
 
     steps:
       - uses: actions/checkout@v2
@@ -32,6 +33,11 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
+          if [ "${{ matrix.pyro-version }}" = "latest" ]; then
+            pip install pyro-ppl
+          else
+            pip install pyro-ppl==${{ matrix.pyro-version }}
+          fi
           pip install .[test]
 
       - name: Lint

chirho/counterfactual/handlers/ambiguity.py

@@ -1,8 +1,10 @@
 import functools
+import typing
 from typing import TypeVar
 
 import pyro
 import pyro.distributions as dist
+import pyro.poutine.indep_messenger
 import torch
 
 from chirho.counterfactual.handlers.selection import (
@@ -23,21 +25,21 @@ class FactualConditioningMessenger(pyro.poutine.messenger.Messenger):
     counterfactual semantics handlers such as :class:`MultiWorldCounterfactual` .
     """
 
-    def _pyro_post_sample(self, msg: dict) -> None:
+    def _pyro_post_sample(self, msg) -> None:
         # expand latent values to include all index plates
         if not msg["is_observed"] and not pyro.poutine.util.site_is_subsample(msg):
             rv, value, event_dim = msg["fn"], msg["value"], len(msg["fn"].event_shape)
             index_plates = get_index_plates()
 
             new_shape = list(value.shape)
             for k in set(indices_of(rv)) - set(indices_of(value, event_dim=event_dim)):
-                dim = index_plates[k].dim
+                dim = typing.cast(int, index_plates[k].dim)
                 new_shape = [1] * ((event_dim - dim) - len(new_shape)) + new_shape
                 new_shape[dim - event_dim] = rv.batch_shape[dim]
 
             msg["value"] = value.expand(tuple(new_shape))
 
-    def _pyro_observe(self, msg: dict) -> None:
+    def _pyro_observe(self, msg) -> None:
         if "name" not in msg["kwargs"]:
             msg["kwargs"]["name"] = msg["name"]
 
@@ -55,7 +57,7 @@ def _dispatched_observe(self, rv, obs: torch.Tensor, name: str) -> torch.Tensor:
     @_dispatched_observe.register(dist.FoldedDistribution)
     @_dispatched_observe.register(dist.Distribution)
     def _observe_dist(
-        self, rv: dist.Distribution, obs: torch.Tensor, name: str
+        self, rv: dist.TorchDistribution, obs: torch.Tensor, name: str
     ) -> torch.Tensor:
         with pyro.poutine.infer_config(config_fn=no_ambiguity):
             with SelectFactual():
@@ -74,7 +76,10 @@ def _observe_dist(
 
     @_dispatched_observe.register
     def _observe_tfmdist(
-        self, rv: dist.TransformedDistribution, value: torch.Tensor, name: str
+        self,
+        rv: dist.TransformedDistribution,
+        value: torch.Tensor,
+        name: str,
     ) -> torch.Tensor:
         tfm = (
             rv.transforms[-1]

chirho/counterfactual/internals.py

@@ -1,9 +1,12 @@
-from typing import Any, Dict
+import typing
+from typing import Any, Mapping, TypedDict
+
+import pyro
 
 from chirho.indexed.ops import indices_of, union
 
 
-def site_is_ambiguous(msg: Dict[str, Any]) -> bool:
+def site_is_ambiguous(msg: Mapping[str, Any]) -> bool:
     """
     Helper function used with :func:`observe` to determine
     whether a site is observed or ambiguous.
@@ -13,17 +16,24 @@ def site_is_ambiguous(msg: Dict[str, Any]) -> bool:
     are fixed/observed (as opposed to random/unobserved).
     """
     rv, obs = msg["args"][:2]
+    infer_dict = msg["infer"]
+    if typing.TYPE_CHECKING:
+        assert infer_dict is not None
     value_indices = indices_of(obs, event_dim=len(rv.event_shape))
     dist_indices = indices_of(rv)
     return (
         bool(union(value_indices, dist_indices)) and value_indices != dist_indices
-    ) or not msg["infer"].get("_specified_conditioning", True)
+    ) or not infer_dict.get("_specified_conditioning", True)
+
+
+class AmbiguityInferDict(pyro.poutine.runtime.InferDict):
+    _specified_conditioning: bool
 
 
-def no_ambiguity(msg: Dict[str, Any]) -> Dict[str, Any]:
+def no_ambiguity(msg: Mapping[str, Any]) -> AmbiguityInferDict:
     """
     Helper function used with :func:`pyro.poutine.infer_config` to inform
     :class:`FactualConditioningMessenger` that all ambiguity in the current
     context has been resolved.
     """
-    return {"_specified_conditioning": True}
+    return AmbiguityInferDict(_specified_conditioning=True)

chirho/dynamical/handlers/trajectory.py

@@ -1,3 +1,4 @@
+import typing
 from typing import Generic, TypeVar
 
 import pyro
@@ -55,7 +56,7 @@ def __init__(self, times: torch.Tensor, is_traced: bool = False):
 
         super().__init__()
 
-    def _pyro_post_simulate(self, msg: dict) -> None:
+    def _pyro_post_simulate(self, msg) -> None:
         initial_state: State[T] = msg["args"][1]
         start_time = msg["args"][2]
 
@@ -74,7 +75,10 @@ def _pyro_post_simulate(self, msg: dict) -> None:
 
         if self.is_traced:
             # This adds the trajectory to the trace so that it can be accessed later.
-            [pyro.deterministic(name, value) for name, value in self.trajectory.items()]
+            [
+                pyro.deterministic(name, typing.cast(torch.Tensor, value))
+                for name, value in self.trajectory.items()
+            ]
 
     def _pyro_simulate_point(self, msg) -> None:
         # Turn a simulate that returns a state into a simulate that returns a trajectory at each of the logging_times
@@ -100,7 +104,9 @@ def _pyro_simulate_point(self, msg) -> None:
 
         # TODO support dim != -1
         idx_name = "__time"
-        name_to_dim = {k: f.dim - 1 for k, f in get_index_plates().items()}
+        name_to_dim = {
+            k: typing.cast(int, f.dim) - 1 for k, f in get_index_plates().items()
+        }
         name_to_dim[idx_name] = -1
 
         if len(timespan) > 2:

chirho/dynamical/internals/_utils.py

@@ -103,7 +103,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             super().__exit__(exc_type, exc_value, traceback)
 
     @typing.final
-    def _process_message(self, msg: Dict[str, Any]) -> None:
+    def _process_message(self, msg) -> None:
         if not self.used and hasattr(self, f"_pyro_{msg['type']}"):
             self.used = True
             super()._process_message(msg)
@@ -119,7 +119,7 @@ def cont(msg: Dict[str, Any]) -> None:
             msg["continuation"] = cont
 
     @typing.final
-    def _postprocess_message(self, msg: Dict[str, Any]) -> None:
+    def _postprocess_message(self, msg) -> None:
         if hasattr(self, f"_pyro_post_{msg['type']}"):
             raise NotImplementedError("ShallowHandler does not support postprocessing")
 

chirho/dynamical/internals/backends/torchdiffeq.py

@@ -1,4 +1,5 @@
 import functools
+import typing
 from typing import Callable, List, Optional, Tuple, TypeVar
 
 import pyro
@@ -151,7 +152,9 @@ def torchdiffeq_simulate_point(
 
     # TODO support dim != -1
     idx_name = "__time"
-    name_to_dim = {k: f.dim - 1 for k, f in get_index_plates().items()}
+    name_to_dim = {
+        k: typing.cast(int, f.dim) - 1 for k, f in get_index_plates().items()
+    }
     name_to_dim[idx_name] = -1
 
     final_idx = IndexSet(**{idx_name: {len(timespan) - 1}})
@@ -249,7 +252,7 @@ def torchdiffeq_simulate_to_interruption(
         dynamics, initial_state, start_time, interruptions, **kwargs
     )
 
-    value = simulate_point(
+    value: State[torch.Tensor] = simulate_point(
         dynamics, initial_state, start_time, interruption_time, **kwargs
     )
     return value, interruption_time, next_interruption

chirho/dynamical/internals/solver.py

@@ -30,7 +30,7 @@ def __init__(
         self.predicate = predicate
         self.callback = callback
 
-    def _pyro_get_new_interruptions(self, msg: dict) -> None:
+    def _pyro_get_new_interruptions(self, msg) -> None:
         if msg["value"] is None:
             msg["value"] = []
         assert isinstance(msg["value"], list)
@@ -42,7 +42,7 @@ def get_new_interruptions() -> List[Interruption]:
     """
     Install the active interruptions into the context.
     """
-    return []
+    return typing.cast(List[Interruption], [])
 
 
 class Solver(Generic[T], pyro.poutine.messenger.Messenger):
@@ -58,11 +58,11 @@ def _prioritize_interruption(h: Interruption[T]) -> Prioritized[Interruption[T]]
             raise NotImplementedError(f"cannot install interruption {h}")
 
     @typing.final
-    def _pyro_simulate(self, msg: dict) -> None:
+    def _pyro_simulate(self, msg) -> None:
         from chirho.dynamical.handlers.interruption import StaticEvent
 
-        dynamics: Dynamics[T] = msg["args"][0]
-        state: State[T] = msg["args"][1]
+        dynamics: Dynamics[T] = typing.cast(Dynamics[T], msg["args"][0])
+        state: State[T] = typing.cast(State[T], msg["args"][1])
         start_time: R = msg["args"][2]
         end_time: R = msg["args"][3]
 
@@ -106,6 +106,7 @@ def _pyro_simulate(self, msg: dict) -> None:
                 if ph.priority > start_time:
                     break
 
+            next_interruption: Optional[Interruption[T]]
             state, start_time, next_interruption = simulate_to_interruption(
                 possible_interruptions,
                 dynamics,
@@ -116,7 +117,8 @@ def _pyro_simulate(self, msg: dict) -> None:
             )
 
             if next_interruption is not None:
-                dynamics, state = next_interruption.callback(dynamics, state)
+                # TODO reenable this check once we figure out why mypy is inferring "Never"s
+                dynamics, state = next_interruption.callback(dynamics, state)  # type: ignore
 
                 for h in possible_interruptions:
                     if h is not next_interruption:

chirho/explainable/handlers/components.py

@@ -1,4 +1,5 @@
 import itertools
+import typing
 from typing import Callable, Iterable, MutableMapping, Optional, TypeVar
 
 import pyro
@@ -91,7 +92,7 @@ def _random_intervention(value: T) -> T:
             support,
             event_shape=event_shape,
         )
-        return pyro.sample(name, proposal_dist)
+        return typing.cast(T, pyro.sample(name, proposal_dist))
 
     return _random_intervention
 
@@ -309,7 +310,7 @@ def _consequent_eq_neq(consequent: T) -> torch.Tensor:
             },
         }
 
-        new_value = scatter_n(
+        new_value: torch.Tensor = scatter_n(
             nec_suff_log_probs_partitioned,
             event_dim=0,
         )
@@ -343,6 +344,6 @@ def __init__(self):
 
         self.supports = {}
 
-    def _pyro_post_sample(self, msg: dict) -> None:
+    def _pyro_post_sample(self, msg) -> None:
         if not pyro.poutine.util.site_is_subsample(msg):
             self.supports[msg["name"]] = msg["fn"].support

chirho/explainable/handlers/explanation.py

@@ -1,4 +1,5 @@
 import contextlib
+import typing
 from typing import Callable, Mapping, Optional, TypeVar, Union
 
 import pyro.distributions.constraints as constraints

chirho/explainable/internals/defaults.py

@@ -2,6 +2,9 @@
 from typing import TypeVar
 
 import pyro
+import pyro.distributions
+import pyro.distributions.constraints
+import pyro.distributions.torch
 import torch
 
 S = TypeVar("S")
@@ -12,7 +15,7 @@
 def uniform_proposal(
     support: pyro.distributions.constraints.Constraint,
     **kwargs,
-) -> pyro.distributions.Distribution:
+) -> pyro.distributions.TorchDistribution:
     """
     This function heuristically constructs a probability distribution over a specified
     support. The choice of distribution depends on the type of support provided.
@@ -44,7 +47,7 @@ def _uniform_proposal_indep(
     *,
     event_shape: torch.Size = torch.Size([]),
     **kwargs,
-) -> pyro.distributions.Distribution:
+) -> pyro.distributions.TorchDistribution:
     d = uniform_proposal(support.base_constraint, event_shape=event_shape, **kwargs)
     return d.expand(event_shape).to_event(support.reinterpreted_batch_ndims)
 
@@ -53,7 +56,7 @@ def _uniform_proposal_indep(
 def _uniform_proposal_integer(
     support: pyro.distributions.constraints.integer_interval,
     **kwargs,
-) -> pyro.distributions.Distribution:
+) -> pyro.distributions.TorchDistribution:
     if support.lower_bound != 0:
         raise NotImplementedError(
             "integer_interval with lower_bound > 0 not yet supported"

chirho/indexed/handlers.py

@@ -124,7 +124,7 @@ def get_mask(
     ) -> torch.Tensor:
         raise NotImplementedError
 
-    def _pyro_sample(self, msg: Dict[str, Any]) -> None:
+    def _pyro_sample(self, msg) -> None:
         if pyro.poutine.util.site_is_subsample(msg):
             return