Raise an error if there is no common scale when model enumerated

[ordabayevy]

It turns out that current subsample scaling logic is not handled correctly in TraceEnum_ELBO (sorry for that). I added some tests to demonstrate that. I think the best place to handle subsampling scaling is in funsor.sum_product.sum_product. For now I added a check for a common scale which raises an error if there is no common scale.

[fehiepsi]

I'm not sure how to check the uniqueness of jax array scales because jax arrays is not hashable. Could you replace this by something like

scale = None
for name in (group_names | group_sum_vars):
    site_scale = model_trace[name]["scale"]
    if isinstance(scale, (int, float)) and isinstance(site_scale, (int, float, type(None))) and (site_scale != scale):
        raise ValueError(...)
    scale = site_scale

Btw, does this mean that we don't support enumeration for models with both global and local variables?

[ordabayevy]

I copied the implementation from Pyro which first checks that array is scalar, converts it to scalar, and then compares them. Does it look alright to you?

https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/traceenum_elbo.py#L37-L41

Btw, does this mean that we don't support enumeration for models with both global and local variables?

Yeah, if you enumerate a global variable than you cannot subsample a local variable that depends on it.
I'm working on it but that will require first changing funsor.sum_product.sum_product to handle plate-wise scaling there.

[fehiepsi]

Because we can't convert a tracer to float, float(site_scale) will fail when scale is a tracer. Maybe we need to raise error if it is a jnp.array?

[ordabayevy]

I see. Changed it so that it raises an error if scale is jnp.array.

[fehiepsi]

LGTM. Thanks, Yerdos.

[ordabayevy]

Thanks for reviewing @fehiepsi .