pyro-ppl / numpyro / 39421073489

93%

push

github

fix(distributions): handle array-bound supports in MaskedDistribution.log_prob (#2380)

Constraint.feasible_like broadcasts the bounds to the shape of the value
it is given. MaskedDistribution.log_prob passed it the value as is, which
can be smaller than the distribution. With array bounds in the base
support this raised:

    Uniform(jnp.array([0.5, 1.0, 2.0]), jnp.array([3.0, 4.0, 5.0])).mask(
        jnp.array([True, False, True])).log_prob(2.2)
    ValueError: Cannot broadcast to shape with fewer dimensions:
    arr_shape=(3,) shape=()

A value of shape (4, 1) raised a broadcasting error as well, where a
Normal base returns shapes (3,) and (4, 3). A scalar observation under
numpyro.handlers.mask failed the same way.

log_prob now broadcasts the value with the distribution shape before
asking for the feasible value. Scalar bounds are unaffected.

1 of 1 new or added line in 1 file covered. (100.0%)

15752 of 16976 relevant lines covered (92.79%)

1.7 hits per line