Type annotate Trace, TraceMessenger, & pyro.poutine.guide
#3299
Conversation
| self, name: str, prior: Distribution | ||
| ) -> Union[Distribution, torch.Tensor]: | ||
| self, name: str, prior: TorchDistributionMixin | ||
| ) -> Union[TorchDistributionMixin, torch.Tensor]: |
There was a problem hiding this comment.
In _pyro_sample above it is expected that posterior has batch_shape. So I changed it here to TorchDistributionMixin because Distribution doesn't have batch_shape. My rule of thumb is that if only sample and log_prob methods are expected is to use Distribution and if additionally batch_shape and event_shape attributes are expected then to use TorchDistributionMixin
ordabayevy
changed the title
Trace, TraceMessenger, & pyro.poutine.guide
pyro/poutine/guide.py
Outdated
| return state | ||
|
|
||
| def __call__(self, *args, **kwargs) -> Dict[str, torch.Tensor]: # type: ignore[override] | ||
| def __call__(self, *args, **kwargs) -> Dict[str, Optional[torch.Tensor]]: # type: ignore[override] |
There was a problem hiding this comment.
I'm surprised this is Optional[Tensor] rather than just Tensor. It might make more sense to assert or filter in the samples loop below, so callers know they are getting tensors.
There was a problem hiding this comment.
Added assert statements in the for loop (a bit longer code but probably better than filtering with if statement in the dict comprehension)
No description provided.