overload

pyro/poutine/block_messenger.py

@@ -82,6 +82,7 @@ def _make_default_hide_fn(
 
 
 def _negate_fn(fn: Callable[[Message], Optional[bool]]) -> Callable[[Message], bool]:
+    # typed version of lambda msg: not fn(msg)
     def negated_fn(msg: Message) -> bool:
         return not fn(msg)
 

pyro/poutine/runtime.py

@@ -14,9 +14,9 @@
     Tuple,
     TypeVar,
     Union,
+    overload,
 )
 
-# overload,
 import torch
 from typing_extensions import ParamSpec, TypedDict
 
@@ -274,18 +274,18 @@ def am_i_wrapped() -> bool:
     return len(_PYRO_STACK) > 0
 
 
-#  @overload
-#  def effectful(
-#      fn: None = ..., type: Optional[str] = ...
-#  ) -> Callable[[Optional[Callable[P, T]]], Callable[P, Union[T, torch.Tensor, None]]]:
-#      ...
-#
-#
-#  @overload
-#  def effectful(
-#      fn: Callable[P, T] = ..., type: Optional[str] = ...
-#  ) -> Callable[P, Union[T, torch.Tensor, None]]:
-#      ...
+@overload
+def effectful(
+    fn: None = ..., type: Optional[str] = ...
+) -> Callable[[Callable[P, T]], Callable[..., Union[T, torch.Tensor, None]]]:
+    ...
+
+
+@overload
+def effectful(
+    fn: Callable[P, T] = ..., type: Optional[str] = ...
+) -> Callable[..., Union[T, torch.Tensor, None]]:
+    ...
 
 
 def effectful(

pyro/primitives.py

@@ -9,11 +9,11 @@
 from typing import Callable, Iterator, Optional, Sequence, Union
 
 import torch
+from torch.distributions import constraints
 
 import pyro.distributions as dist
 import pyro.infer as infer
 import pyro.poutine as poutine
-from pyro.distributions import constraints
 from pyro.params import param_with_module_name
 from pyro.params.param_store import ParamStoreDict
 from pyro.poutine.plate_messenger import PlateMessenger
@@ -48,9 +48,7 @@ def clear_param_store() -> None:
     _PYRO_PARAM_STORE.clear()
 
 
-_param: Callable[..., torch.Tensor] = effectful(
-    _PYRO_PARAM_STORE.get_param, type="param"
-)
+_param = effectful(_PYRO_PARAM_STORE.get_param, type="param")
 
 
 def param(
@@ -84,9 +82,11 @@ def param(
     :rtype: torch.Tensor
     """
     # Note effectful(-) requires the double passing of name below.
-    return _param(
+    value = _param(
         name, init_tensor, constraint=constraint, event_dim=event_dim, name=name
     )
+    assert isinstance(value, torch.Tensor)
+    return value
 
 
 def _masked_observe(