handlers & py.typed (#3321)

pyro-ppl · Feb 6, 2024 · 6d2a56f · 6d2a56f
1 parent df34383
commit 6d2a56f
Show file tree

Hide file tree

Showing 5 changed files with 380 additions and 63 deletions.
diff --git a/pyro/poutine/collapse_messenger.py b/pyro/poutine/collapse_messenger.py
@@ -3,7 +3,7 @@
 
 
 from functools import reduce, singledispatch
-from typing import TYPE_CHECKING, FrozenSet, Tuple
+from typing import TYPE_CHECKING, Any, FrozenSet, Tuple
 
 from typing_extensions import Self
 
@@ -85,7 +85,7 @@ class CollapseMessenger(TraceMessenger):
 
     _coerce = None
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         if CollapseMessenger._coerce is None:
             import funsor
             from funsor.distribution import CoerceDistributionToFunsor