Support IndependentConstraint in Distribution._infer_param_domain

[eb8680]

Addresses #386. Blocking all other PRs.

This PR updates the logic in Distribution._infer_param_domain to handle the new _IndependentConstraint in NumPyro added in pyro-ppl/numpyro#876, which should fix the CI failures in #430 and other PRs.

Tested:

• Exercised by existing tests that are currently failing, including test_distributions_generic::test_generic_sample for CategoricalLogits

[fehiepsi]

Thanks for the fix, Eli!

@fritzo I guess we can add

if isinstance(base_constraint, _IndependentConstraint):
    batch_ndims = base_constraint.batch_ndims + batch_ndims
    base_constraint = base_constraint.base_constraint

to simplify the while loop here. What do you think?

[fritzo]

Note the constraint names will probably be made public in the future so we should do something like

support_name = type(support).__name__.lstrip("_")

and reference names like IndependentConstraint and Real, so as to be forwards-compatible, and to be compatible with both Pyro and NumPyro and future versions of PyTorch.

[fritzo]

Can you change this to

support_name = type(support).__name__.lstrip("_")

event_dim = 0
while support_name == "IndependentConstraint":
    ...

so as to be compatible with all backends? Note in Pyro the class is named IndependentConstraint; in PyTorch the class will initially be named _IndependentConstraint; and in a future PyTorch the class will be renamed to IndependentConstraint.

[eb8680]

We could also duplicate _infer_param_domain in each backend so that dispatch happens on constraint types rather than names.

[fehiepsi]

and in a future PyTorch the class will be renamed to IndependentConstraint

Does it apply to other constraints as well?

[fritzo]

yes, see plans in pytorch/pytorch#50616