Use ExpandedDistribution in numpyro.plate

[neerajprad]

Since #574, I think we should be able to switch our existing logic to use ExpandedDistribution instead of our current logic. This will align the underlying implementation with that of Pyro and possibly aid with funsor integration.

[fehiepsi]

I guess we can also use MaskedDistribution for mask handler to simplify mask logic in log_density computation. WDYT?

Edit: I think we should do so. Currently, unconstrained_reparam does not work correctly under mask handler because we put log_det term under mask handler, which is wrong. :(

[neerajprad]

I guess we can also use MaskedDistribution for mask handler to simplify mask logic in log_density computation. WDYT?

For ExpandedDistribution my only use case is supporting funsor integration, but if we don't need it for that, I'm happy to keep things simple as they are now. What was your reason for migrating it to NumPyro?

Likewise, I would say the same for MaskedDistribution, but it seems that you already have a use case for that, so let us use MaskedDistribution for that instead.