pyro-ppl / numpyro / 39302668734

93%

Build Type

push

github

Committed by web-flow

Commit Message

refactor(distributions): derive the interval supports from low and high (#2376)

* refactor(distributions): derive the interval supports from low and high

Uniform, LogUniform and DiscreteUniform kept a copy of their support,
built from the bounds in __init__ and carried as an extra pytree field,
so the bounds were flattened twice and the copy fell behind when only
the parameter leaves changed: after a gradient step moved
Uniform(1.0, 3.0) to (1.5, 3.5), support still accepted 1.2 and
log_prob(3.2) was -inf under validate_args.

The support is now computed from low and high on access, as for Pareto
and Binomial. Uniform and DiscreteUniform flatten to their two
parameters and LogUniform to four, because it still carries its
log-space base_dist, which log_prob and sample keep using. The vmap_over
lines that only synced the copy are removed.

* test: check only the support after a leaf update, drop the leaf count test

LogUniform.log_prob still reads base_dist, which keeps the old bounds
after a leaf update, so the log_prob assertion passed on a value that
did not follow the new bounds. Check only the support for both classes.

Drop test_interval_dist_pytree_leaves: it pins how the classes flatten,
not what the support does.

Coverage Stats

6 of 6 new or added lines in 2 files covered. (100.0%)

4 existing lines in 1 file now uncovered.

15752 of 16976 relevant lines covered (92.79%)

1.7 hits per line