diff --git a/pyro/nn/__init__.py b/pyro/nn/__init__.py index e55e7356f6..310805a7c2 100644 --- a/pyro/nn/__init__.py +++ b/pyro/nn/__init__.py @@ -14,6 +14,7 @@ PyroModuleList, PyroParam, PyroSample, + PyroSamplePlateScope, pyro_method, ) @@ -28,4 +29,5 @@ "PyroSample", "pyro_method", "PyroModuleList", + "PyroSamplePlateScope", ] diff --git a/pyro/nn/module.py b/pyro/nn/module.py index b84a17875b..d3eff1d823 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -37,6 +37,7 @@ def _copy_to_script_wrapper(fn): Any, Callable, Dict, + Iterable, Iterator, List, NamedTuple, @@ -53,8 +54,10 @@ def _copy_to_script_wrapper(fn): import pyro import pyro.params.param_store +from pyro.distributions.torch_distribution import TorchDistributionMixin from pyro.ops.provenance import detach_provenance -from pyro.poutine.runtime import _PYRO_PARAM_STORE +from pyro.poutine.messenger import Messenger +from pyro.poutine.runtime import _PYRO_PARAM_STORE, InferDict _MODULE_LOCAL_PARAMS: bool = False @@ -63,7 +66,6 @@ def _copy_to_script_wrapper(fn): _PyroModule = TypeVar("_PyroModule", bound="PyroModule") if TYPE_CHECKING: - from pyro.distributions.torch_distribution import TorchDistributionMixin from pyro.params.param_store import StateDict @@ -232,6 +234,48 @@ def __get__( return value +class _PyroSampleInferDict(InferDict): + _original_pyrosample_dist: TorchDistributionMixin + + +class PyroSamplePlateScope(Messenger): + """ + Handler for executing PyroSample statements in a more intuitive plate context. + """ + + def __init__(self, allowed_plates: Iterable[str] = ()): + self._inner_allowed_plates = frozenset(allowed_plates) + + def __enter__(self): + self._plates = ( + frozenset(p.name for p in pyro.poutine.runtime.get_plates()) + | self._inner_allowed_plates + ) + return super().__enter__() + + def _is_local_plate(self, m: Messenger) -> bool: + return ( + isinstance(m, pyro.poutine.plate_messenger.PlateMessenger) + and m.name not in self._plates + ) + + def _pyro_sample(self, msg) -> None: + if not msg["infer"].get("_original_pyrosample_dist", None): + return + msg["stop"] = True + msg["done"] = True + with pyro.poutine.messenger.block_messengers( + lambda m: m is self or self._is_local_plate(m) + ): + d = msg["infer"].pop("_original_pyrosample_dist") + msg["value"] = pyro.sample( + msg["name"], + d, + obs=msg["value"] if msg["is_observed"] else None, + infer=msg["infer"], + ) + + def _make_name(prefix: str, name: str) -> str: return "{}.{}".format(prefix, name) if prefix else name @@ -615,7 +659,13 @@ def __getattr__(self, name: str) -> Any: value = ( pyro.deterministic(fullname, prior) if isinstance(prior, torch.Tensor) - else pyro.sample(fullname, prior) + else pyro.sample( + fullname, + prior, + infer=_PyroSampleInferDict( + _original_pyrosample_dist=prior + ), + ) ) context.set(fullname, value) return value diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 07c4daedd1..665a63d5a5 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -1084,3 +1084,46 @@ def forward(self): with pyro.settings.context(module_local_params=use_module_local_params): model = Model() pyro.render_model(model) + + +def test_pyrosample_platescope(): + + class Model(pyro.nn.PyroModule): + def __init__(self, num_inputs, num_outputs): + super().__init__() + self.num_inputs = num_inputs + self.num_outputs = num_outputs + self.linear = pyro.nn.PyroModule[torch.nn.Linear](num_inputs, num_outputs) + self.linear.weight = pyro.nn.PyroSample( + dist.Normal(0, 1).expand([num_outputs, num_inputs]).to_event(2) + ) + self.linear.bias = pyro.nn.PyroSample( + dist.Normal(0, 1).expand([num_outputs]).to_event(1) + ) + + @pyro.nn.PyroSample + def scale(self): + return ( + pyro.distributions.LogNormal(0, 1) + .expand([self.num_outputs]) + .to_event(1) + ) + + @pyro.nn.PyroSamplePlateScope() + def forward(self, x): + with pyro.plate("data", x.shape[-2], dim=-1): + assert ( + len(self.linear.weight.shape) == 2 + or self.linear.weight.shape[-3] != 1 + ) # sampled outside data plate + loc = self.linear(x) + assert ( + len(self.scale.shape) == 1 or self.scale.shape[-2] == 1 + ) # sampled outside data plate + y = pyro.sample("y", dist.Normal(loc, self.scale).to_event(1)) + assert y.shape[-2] == x.shape[-2] # ordinary pyro.sample statement + return y + + model = Model(3, 2) + x = torch.randn(4, 3) + assert model(x).shape == (4, 2)