Fixes pyro-ppl/pyro#3419 which is actually a `torch` bug that can be replicated by the below code: ``` from torch import rand from torch.distributions import MixtureSameFamily, Categorical, Binomial max_count = 20 probs = rand(10, 5) binom_probs = rand(10, 5) d = MixtureSameFamily(Categorical(probs=probs), Binomial(max_count, binom_probs)) d.log_prob(d.sample()) ``` which results in: ``` Traceback (most recent call last): File "test.py", line 11, in <module> d.log_prob(d.sample()) File "pytorch\torch\distributions\mixture_same_family.py", line 168, in log_prob self._validate_sample(x) File "pytorch\torch\distributions\distribution.py", line 315, in _validate_sample valid = support.check(value) ^^^^^^^^^^^^^^^^^^^^ File "pytorch\torch\distributions\constraints.py", line 307, in check (value % 1 == 0) & (self.lower_bound <= value) & (value <= self.upper_bound) ^^^^^^^^^^^^^^^^^^^^^^^^^ RuntimeError: The size of tensor a (10) must match the size of tensor b (5) at non-singleton dimension 1 ``` ### Fix explanation (only for cases when the component distribution contains parameters with batch dimenisons) - The failure is due to sample validation taking place before padding in `MixtureSameFamily.log_prob`, and hence the fix is to pad before doing sample validation. - The fix itself does not alter the calculations at all. It only affects the sample validation process. - The failure does not occur with the component distribution set to the `Normal` distribution, as its validation is not defined elementwise (the validation itself is elementwise). - I've split the `test_mixture_same_family_log_prob` test into two tests based on the `Normal` and `Binomial` distributions. - Initially, the `Binomial` version of the test did not fail, but this was due to the component distribution having equal batch dimensions of (5, 5) so I changed it to (10, 5). ### Updated fix explanation (for all cases) - The previous fix caused a bug in sample shape validation (which is done correctly) due to the padding taking place before the sample validation. - The updated fix corrects the support to reflect the fact that the support of `MixtureSameFamily` is equal to the support of its components distribution with the first event dimension removed. - This issue was already anticipated in the [code](https://github.com/pytorch/pytorch/blob/331423e5c24170b218e743b3392acbad4480340d/torch/distributions/mixture_same_family.py#L127). Pull Request resolved: #151317 Approved by: https://github.com/albanD, https://github.com/fritzo

Uh oh!

[bug] Batched mixture of binomials fails validation #3419

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions