Fix EqualizeMessenger type annotations.

pyro-ppl · Sep 23, 2024 · d3e46b1 · d3e46b1
1 parent b3c7851
commit d3e46b1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pyro/poutine/handlers.py b/pyro/poutine/handlers.py
@@ -306,14 +306,16 @@ def escape(  # type: ignore[empty-body]
 def equalize(
     sites: Union[str, List[str]],
     type: Optional[str],
-) -> ConditionMessenger: ...
+    keep_dist: Optional[bool],
+) -> EqualizeMessenger: ...
 
 
 @overload
 def equalize(
     fn: Callable[_P, _T],
     sites: Union[str, List[str]],
     type: Optional[str],
+    keep_dist: Optional[bool],
 ) -> Callable[_P, _T]: ...
 
 
@@ -322,6 +324,7 @@ def equalize(  # type: ignore[empty-body]
     fn: Callable[_P, _T],
     sites: Union[str, List[str]],
     type: Optional[str],
+    keep_dist: Optional[bool],
 ) -> Union[EqualizeMessenger, Callable[_P, _T]]: ...