Properly handle contraction of guide plates in TraceEnum_ELBO
#1537
Conversation
| def model(params): | ||
| with pyro.plate("a_axis", size=2): | ||
| a = pyro.sample("a", dist.Categorical(jnp.array([0.2, 0.8]))) | ||
| pyro.sample("b", dist.Normal(jnp.sum(a), 1.0), obs=1) |
There was a problem hiding this comment.
It seems to me that this invalidates this restriction: https://pyro.ai/examples/enumeration.html#Restriction-2:-no-downstream-coupling? Could you elaborate why we can enumerate here?
There was a problem hiding this comment.
It's not being enumerated here but being used as a non-reparameterizable site. I should have used other distribution like Poisson to make it less confusing :)
But the point is that b depends on a non-reparametrizable site a which has a_axis plate. The dice_factor needs to be product-contracted to eliminate the extra a_axis plate before multiplying the cost term for site b. Instead a_axis is being passed to _eager_contract_tensors as reduced_vars and sum-contracted with logsumexp. Hope this clarifies it.
There was a problem hiding this comment.
I see. So we are using TraceEnum_ELBO but enumeration is disabled for those cases. Could you add a warning for this?
There was a problem hiding this comment.
Do you mean in general or for this test?
There was a problem hiding this comment.
I mean in general. I think using TraceEnum_ELBO without enumeration is confusing. Maybe raise error if we can't enumerate sites with infer={"enumerate": "parallel"}?
There was a problem hiding this comment.
I changed the distribution to Poisson in the test since it can be any non-reparameterizable distributions. I can open another issue/PR for enumeration configuration since it is a separate issue.
There was a problem hiding this comment.
sounds reasonable to me, thanks!
|
Thanks for reviewing @fehiepsi ! |
There was a bug in
TraceEnum_ELBOdemonstrated in a newtest_enum_elbo::test_guide_plate_contractiontest. When a cost term depends on a non-reparameterizable guide site with extra plate dims they need to be product-contracted. Instead these plate dims were passed to_eager_contract_tensorsasreduced_varsand sum-contracted. So I replaced_eager_contract_tensorswithfunsor.sum_product.sum_productwhich can do plated sum-products and eliminate plates.