Enabling sampling with intermediates for `ExpandedDistribution` (#909)
Enabling sampling with intermediates for `ExpandedDistribution` (#909)

* Simplifying ExpandedDistribution._sample

It involved some elaborate shuffling in sample/batch dimensions
that is unneccessary (due to independence in these dimensions)
and didn't mitigate the need for a reshape in the end,
which achieves the same results all on its own.

* Added sample_with_intermediates for ExpandedDistribution

Assuming that intermediates are batched similarly to direct samples.
pyro-ppl/numpyro