Add Distribution.infer_shapes() for static shape analysis

[fritzo]

Addresses pyro-ppl/funsor#412
Ports .infer_shape() from pyro-ppl/pyro#2739
Ports dependent_property from pytorch/pytorch#50581

This adds a static method Disribution.infer_shapes() to statically infer (batch_shape, event_shape) from the shapes of distribution parameters. This should speed up funsor computations as described in pyro-ppl/funsor#412. Note this works for only a subset of distributions; shapes cannot currently be statically inferred for distributions with non-tensor inputs such as TransformedDistribution which takes a distribution, or Delta which takes an event_dim: int.

This generic logic requires static .event_dim attributes on dependent supports. Because NumPyro uses Constraint.__call__() rather than Constraint.check() I've added a type check to determine behavior in _Dependent.__call__().

Tested

• added a unit test to test_distributions.py

[fritzo]

@fehiepsi any idea why tests are failing? The failures all seem to be numerical precision issues, which is weird because this PR only modifies type metadata 😕

EDIT Oops, it was probably my incorrect LogNormal.support 😊

[fehiepsi]

Looks great to me! Thanks, @fritzo!

[fehiepsi]

I'm not sure what is a good solution here...

[fritzo]

I guess to be completely correct we'd want something like Uniform.arg_constraints:

arg_constraints = {'v': constraints.dependent(is_discrete=False),
                   'log_density': constraints.real}

but I was worried that might break something and seemed out of scope for this PR.

[fritzo]

Here's a possible fix: https://github.com/pyro-ppl/numpyro/compare/fix-delta-constraint

[fehiepsi]

It is surprising to me that this works for @constraints.dependent_property. I have been thinking that we can't make infer_shapes a classmethod for many distributions.

[fritzo]

Yeah this is a kind of cool trick. The two versions

self.support.event_dim
type(self).support.event_dim

follow entirely different code paths: the former runs the property method to compute a Constraint instance then extracts the .event_dim of that instance. The latter cannot create a Constraint instance but instead extracts a class attribute of the dependent_property itself. In PyTorch I've added tests that these two code paths agree.

[fehiepsi]

Totally make sense to me! Thanks for explaning and being creative!

[fritzo]

Thanks for reviewing @fehiepsi !