pyro-ppl / numpyro / 32653608150

92%

push

github

Drop the masked logsumexp reduction in mixture `log_prob` (#2250)

The where= kwarg of jax.nn.logsumexp in _MixtureBase.log_prob forces a
masked reduction that defeats XLA's fused reduction path (3x slower for
MixtureSameFamily log_prob at 10k values on CPU). The mask is redundant:
the value-preserving jnp.where introduced in #1874 already yields the
same values and stops gradients into -inf component log-probs; the
where= kwarg was added incidentally in the type-hints PR #2032. Values
and gradients are unchanged, including at zero mixing weights and for
all-components -inf rows; add regression tests for both cases.

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

14631 of 15839 relevant lines covered (92.37%)

1.69 hits per line