pyro-ppl / numpyro / 39421793533

93%

push

github

fix(distributions): build the censored feasible values at the batch shape (#2379)

Constraint.feasible_like broadcasts the bounds to the shape of the value
it is given. LeftCensoredDistribution, RightCensoredDistribution and
IntervalCensoredDistribution passed it a scalar in __init__, and
IntervalCensoredDistribution.log_prob passed it x1, which can be smaller
than the batch. 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]))
    LeftCensoredDistribution(base, censored, validate_args=True)
    ValueError: Cannot broadcast to shape with fewer dimensions:
    arr_shape=(3,) shape=()

while validate_args=False built the same distribution. The same bounds
made log_prob(jnp.array([2.1, 2.5])) on the interval class raise the
same error, and a (4, 1, 2) value raise a broadcasting error, where a
Normal base returns shapes (3,) and (4, 3).

__init__ now passes zeros of the base shape, and log_prob broadcasts x1
to its own shape and the batch shape before asking for the feasible
value. Scalar bounds are unaffected.

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

15760 of 16976 relevant lines covered (92.84%)

1.7 hits per line