Scale factors across plate dims in partial_sum_product

[ordabayevy]

One of the features not supported by TraceEnum_ELBO is that you cannot subsample a local variable when it depends on a global variable that is enumerated in the model because it requires a common scale:

@config_enumerate
def model(data):
    # Global variables.
    locs = torch.tensor([1., 10.])
    assignment = pyro.sample('assignment', dist.Categorical(torch.ones(2)))
    with pyro.plate('data', len(data), subsample_size=2) as ind:  # cannot subsample here
        # Local variables.
        pyro.sample('obs', dist.Normal(locs[assignment], 1.), obs=data[ind])

def guide(data):
    pass

This has been asked on the forum as well: https://forum.pyro.ai/t/enumeration-and-subsampling-expected-all-enumerated-sample-sites-to-share-common-poutine-scale/4938

A solution I'm proposing here is to perform plate-wise scaling inside the partial_sum_product by passing the plate_to_scale dictionary. Then whenever a plate is reduced we can scale the factor:

# inside partial_sum_product
f = f.reduce(prod_op, reduced_plates)
f_scales = [plate_to_scale[plate] for plate in reduced_plates if plate in plate_to_scale]
if f_scales:
    scale = reduce(ops.mul, f_scales)
    f = pow_op(f, scale)

• Added tests

[ordabayevy]

@fritzo @eb8680 can you have a look at this? Do changes to partial_sum_product look reasonable to you?

[ordabayevy]

@eb8680 I opened a PR in NumPyro (pyro-ppl/numpyro#1572) where I tried to expand the math and check whether expectations match for a full- and a mini-batch log-likelihoods. Can you verify that the math is correct?

[fritzo]

This should be NotImplementedError

[fritzo]

Move this out to a PROD_TO_POWER dict or similar, see ops

[fritzo]

Could we wrap this in an if plate_to_scale: guard to improve readability?

[fritzo]

ditto: if plate_to_scale: ...

[fritzo]

default to None, maybe comment on datatype

[ordabayevy]

@fritzo can you review the changes please? Also I believe the rule for jax tests needs to be changed in the Settings to use python 3.9 instead of 3.8. Jax doesn't support 3.8 anymore.

[fritzo]

LGTM, thanks for addressing my comments.