refactor(infer.elbo): add type hints to elbo module (#2028)

* initial exploration of typing for infer.elbo

* continue adding types, cleanup code dead ends

* add _typing module, add ANN rule, finish elbo module

* refactor(infer.elbo): bind paramspecs to methods, add type aliases

* refactor(infer.elbo): ParticleT -> LossT and update trace elbo annotations

* fix: import TypeAlias from typing_extensions

* exclude python 3.9 from lint/typecheck, MessageT -> Message

Author: brendancooley

.github/workflows/ci.yml

@@ -36,6 +36,7 @@ jobs:
         pip install -r docs/requirements.txt
         pip freeze
     - name: Lint with mypy and ruff
+      if: matrix.python-version != '3.9'
       run: |
         make lint
     - name: Build documentation

numpyro/_typing.py

@@ -0,0 +1,14 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import OrderedDict
+from collections.abc import Callable
+from typing import Any
+
+from typing_extensions import ParamSpec, TypeAlias
+
+P = ParamSpec("P")
+ModelT: TypeAlias = Callable[P, Any]
+
+Message: TypeAlias = dict[str, Any]
+TraceT: TypeAlias = OrderedDict[str, Message]

numpyro/contrib/stochastic_support/dcc.py

@@ -8,7 +8,6 @@
 import jax
 from jax import random
 import jax.numpy as jnp
-from jax.typing import ArrayLike
 
 import numpyro.distributions as dist
 from numpyro.handlers import condition, seed, trace
@@ -61,7 +60,7 @@ def __init__(self, model: Callable, num_slp_samples: int, max_slps: int) -> None
         self.max_slps: int = max_slps
 
     def _find_slps(
-        self, rng_key: ArrayLike, *args: Any, **kwargs: Any
+        self, rng_key: jax.Array, *args: Any, **kwargs: Any
     ) -> dict[str, OrderedDictType]:
         """
         Discover the straight-line programs (SLPs) in the model by sampling from the prior.
@@ -109,7 +108,7 @@ def _get_branching_trace(self, tr: dict[str, Any]) -> OrderedDictType:
     @abstractmethod
     def _run_inference(
         self,
-        rng_key: ArrayLike,
+        rng_key: jax.Array,
         branching_trace: OrderedDictType,
         *args: Any,
         **kwargs: Any,
@@ -119,7 +118,7 @@ def _run_inference(
     @abstractmethod
     def _combine_inferences(
         self,
-        rng_key: ArrayLike,
+        rng_key: jax.Array,
         inferences: dict[str, Any],
         branching_traces: dict[str, OrderedDictType],
         *args: Any,
@@ -128,7 +127,7 @@ def _combine_inferences(
         raise NotImplementedError
 
     def run(
-        self, rng_key: ArrayLike, *args: Any, **kwargs: Any
+        self, rng_key: jax.Array, *args: Any, **kwargs: Any
     ) -> Union[DCCResult, SDVIResult]:
         """
         Run inference on each SLP separately and combine the results.
@@ -209,7 +208,7 @@ def __init__(
 
     def _run_inference(
         self,
-        rng_key: ArrayLike,
+        rng_key: jax.Array,
         branching_trace: OrderedDictType,
         *args: Any,
         **kwargs: Any,
@@ -226,7 +225,7 @@ def _run_inference(
 
     def _combine_inferences(  # type: ignore[override]
         self,
-        rng_key: ArrayLike,
+        rng_key: jax.Array,
         samples: dict[str, Any],
         branching_traces: dict[str, OrderedDictType],
         *args: Any,
@@ -244,7 +243,7 @@ def _combine_inferences(  # type: ignore[override]
         """
 
         def log_weight(
-            rng_key: ArrayLike,
+            rng_key: jax.Array,
             i: int,
             slp_model: Callable,
             slp_samples: dict[str, Any],

numpyro/contrib/stochastic_support/sdvi.py

@@ -5,7 +5,6 @@
 
 import jax
 import jax.numpy as jnp
-from jax.typing import ArrayLike
 
 from numpyro.contrib.stochastic_support.dcc import (
     RunInferenceResult,
@@ -98,7 +97,7 @@ def __init__(
 
     def _run_inference(
         self,
-        rng_key: ArrayLike,
+        rng_key: jax.Array,
         branching_trace: OrderedDictType,
         *args: Any,
         **kwargs: Any,
@@ -120,17 +119,17 @@ def _run_inference(
 
     def _combine_inferences(  # type: ignore[override]
         self,
-        rng_key: ArrayLike,
+        rng_key: jax.Array,
         guides: dict[str, tuple[AutoGuide, dict[str, Any]]],
         branching_traces: dict[str, OrderedDictType],
         *args: Any,
         **kwargs: Any,
     ) -> SDVIResult:
         """Weight each SLP proportional to its estimated ELBO."""
-        elbos = {}
+        elbos: dict[str, jax.Array] = {}
         for bt, (guide, param_map) in guides.items():
             slp_model = condition(self.model, branching_traces[bt])
-            elbos[bt] = -Trace_ELBO(num_particles=self.combine_elbo_particles).loss(
+            elbos[bt] = -Trace_ELBO(num_particles=self.combine_elbo_particles).loss(  # type: ignore
                 rng_key, param_map, slp_model, guide, *args, **kwargs
             )
 

numpyro/handlers.py

@@ -91,6 +91,8 @@ def seeded_model(data):
    -874.89813
 """
 
+from __future__ import annotations
+
 from collections import OrderedDict
 from types import TracebackType
 from typing import Callable, Optional, Union
@@ -103,12 +105,12 @@ def seeded_model(data):
 from jax.typing import ArrayLike
 
 import numpyro
+from numpyro._typing import Message, TraceT
 from numpyro.distributions.distribution import COERCIONS
 from numpyro.primitives import (
     _PYRO_STACK,
     CondIndepStackFrame,
     DistributionLike,
-    Message,
     Messenger,
     apply_stack,
     plate,
@@ -163,9 +165,9 @@ class trace(Messenger):
                       'value': Array(-0.20584235, dtype=float32)})])
     """
 
-    def __enter__(self) -> OrderedDict[str, Message]:  # type: ignore [override]
+    def __enter__(self) -> TraceT:  # type: ignore [override]
         super(trace, self).__enter__()
-        self.trace: OrderedDict[str, Message] = OrderedDict()
+        self.trace: TraceT = OrderedDict()
         return self.trace
 
     def postprocess_message(self, msg: Message) -> None:
@@ -180,7 +182,7 @@ def postprocess_message(self, msg: Message) -> None:
         )
         self.trace[msg["name"]] = msg.copy()
 
-    def get_trace(self, *args, **kwargs) -> OrderedDict[str, Message]:
+    def get_trace(self, *args, **kwargs) -> TraceT:
         """
         Run the wrapped callable and return the recorded trace.
 
@@ -225,7 +227,7 @@ class replay(Messenger):
     def __init__(
         self,
         fn: Optional[Callable] = None,
-        trace: Optional[OrderedDict[str, Message]] = None,
+        trace: Optional[TraceT] = None,
     ) -> None:
         assert trace is not None
         self.trace = trace
@@ -357,7 +359,7 @@ def process_message(self, msg: Message) -> None:
             if isinstance(msg["fn"], Funsor) or isinstance(msg["value"], (str, Funsor)):
                 msg["stop"] = True
 
-    def __enter__(self) -> OrderedDict[str, Message]:  # type: ignore [override]
+    def __enter__(self) -> TraceT:  # type: ignore [override]
         self.preserved_plates = frozenset(
             h.name for h in _PYRO_STACK if isinstance(h, plate)
         )
@@ -451,7 +453,7 @@ def __init__(
             raise ValueError("Only one of `data` or `condition_fn` should be provided.")
         super(condition, self).__init__(fn)
 
-    def process_message(self, msg):
+    def process_message(self, msg: Message) -> None:
         if (msg["type"] != "sample") or msg.get("_control_flow_done", False):
             if msg["type"] == "control_flow":
                 if self.data is not None:
@@ -465,6 +467,7 @@ def process_message(self, msg):
         if self.data is not None:
             value = self.data.get(msg["name"])
         else:
+            assert self.condition_fn is not None
             value = self.condition_fn(msg)
 
         if value is not None:
@@ -804,9 +807,9 @@ class seed(Messenger):
 
     def __init__(
         self,
-        fn: Optional[Callable] = None,
-        rng_seed: Optional[Array] = None,
-        hide_types: Optional[list[str]] = None,
+        fn: Callable | None = None,
+        rng_seed: Array | int | None = None,
+        hide_types: list[str] | None = None,
     ) -> None:
         if rng_seed is not None:
             if not is_prng_key(rng_seed) and (