Add FunsorDistribution to wrap funsors for use in Pyro #170
Conversation
| return value | ||
|
|
||
| def expand(self, batch_shape, _instance=None): | ||
| raise NotImplementedError("TODO") |
There was a problem hiding this comment.
.expand will be necessary for most Pyro use cases, right? Are you planning to implement this generically in a followup PR, or case-by-case for specific FunsorDistribution subclasses? Also, can we use funsor.Independent to implement FunsorDistribution.to_event?
There was a problem hiding this comment.
I'm planning to implement .expand() generically when I eventually need it.
Hmm yes, I think we might be able to use funsor.Independent for .to_event().
| return tensor | ||
|
|
||
|
|
||
| def dist_to_funsor(pyro_dist, reinterpreted_batch_ndims=0): |
There was a problem hiding this comment.
We should be able to implement dist_to_funsor generically and use it to replace the current funsor.distributions as suggested in #159, with a few extra eager patterns for Categorical, Normal/MultivariateNormal, and Delta
There was a problem hiding this comment.
In theory, yes. In practice, there is a mismatch between the dtype and shapes in funsor versus distributions. I welcome any effort to implement this generically, but I suspect every distribution will be a new edge case.
|
Thanks for reviewing! |
Addresses #177
This adds a basic
FunsorDistributionclass inheriting fromTorchDistribution. The class is intended to wrap aFunsoras aDistributionobject for use in Pyro.I've created a new module
funsor.pyroand I plan to create new submodules likefunsor.pyro.hmm#171 containing various hidden-markov models (discrete, gaussian, slds, ...). Each of these distributions will use fancy funsor syntax internally, but will expose their inference algorithms as black-boxDistributionswith.sample()and.log_prob()methods. These distributions should be fully compatible with broadcasting, plates, and enumeration dims (which should pass right through these algorithms).Tested
Categoricaldistribution from a funsor