Add more tests for collapse handler

[fehiepsi]

Originally, this is used to test for funsor PR: pyro-ppl/funsor#392. But I found that we need more tests when updating the code to match Pyro version.

To pass some tests, I need to eagerly expand an ExpandedDistribution using _eager_expand_fn so that: Normal(0, 1).expand([10]) is still a Normal instance. This is not needed in Pyro because fn.expand(...) returns a new instance of type(fn).

I added some failing tests for discussions. Addressing them can be deferred to a follow-up PR in funsor.

[fritzo]

Do I understand correctly that this function eliminates lazy ExpandedDistributions and replaces them with eagerly tree_map-expanded distributions, and that this is needed because Funsor distributions do not support Expanded distributions? If so could you add a commend explaining that?

I guess an alternative would be to support ExpandedDistributions in Funsor, or to do this conversion in the funsor layer. @eb8680 does that seem possible?

[fehiepsi]

Yes, it is. I'll add a comment for clarity. I thought from our last discussion, this stuff should be treated in Pyro/NumPyro?

[eb8680]

I guess an alternative would be to support ExpandedDistributions in Funsor, or to do this conversion in the funsor layer.

We could certainly implement an equivalent of .expand. I think the easiest thing to do to address this case would be to add a to_funsor pattern for ExpandedDistribution.

[fehiepsi]

add a to_funsor pattern for ExpandedDistribution

Yeah, it seems that you both agree to do this. I can take an effort for it.

[eb8680]

I think it would basically be the same as what you have here, plus a final step that calls to_funsor on the eagerly expanded result.