Drop the masked logsumexp reduction in mixture log_prob

numpyro/distributions/mixtures.py

@@ -168,14 +168,12 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
     def log_prob(self, value: ArrayLike, intermediates=None) -> ArrayLike:
         del intermediates
         sum_log_probs = self.component_log_probs(value)
+        # the value-preserving where stops gradients from flowing into -inf
+        # component log-probs (e.g. at zero mixing weights); see #1874
         safe_sum_log_probs = jnp.where(
             jnp.isneginf(sum_log_probs), -jnp.inf, sum_log_probs
         )
-        return jax.nn.logsumexp(
-            safe_sum_log_probs,
-            where=~jnp.isneginf(sum_log_probs),  # for numerical stability
-            axis=-1,
-        )
+        return jax.nn.logsumexp(safe_sum_log_probs, axis=-1)
 
 
 class MixtureSameFamily(_MixtureBase):

test/test_distributions_mixture.py

@@ -203,3 +203,35 @@ def test_mixture_rejects_parameter_dependent_components(component_dist):
 def test_mixture_accepts_parameter_free_components(component_dist):
     mixing_dist = dist.Categorical(probs=np.array([0.3, 0.7]))
     dist.MixtureSameFamily(mixing_dist, component_dist)
+
+
+def test_mixture_log_prob_grad_at_zero_weights():
+    # regression test for #1874: gradients must not become NaN when some
+    # mixing weights are exactly zero
+    def loss(probs):
+        mixture = dist.MixtureSameFamily(
+            dist.Categorical(probs=probs),
+            dist.Normal(jnp.arange(3.0), jnp.ones(3)),
+        )
+        return mixture.log_prob(jnp.array([0.5, 2.0])).sum()
+
+    grad = jax.grad(loss)(jnp.array([0.5, 0.5, 0.0]))
+    assert not jnp.isnan(grad).any()
+
+
+def test_mixture_log_prob_grad_all_components_neg_inf():
+    # if every component assigns -inf to a value, the log_prob is -inf and
+    # gradients must not become NaN
+    mixture = dist.MixtureGeneral(
+        dist.Categorical(probs=jnp.array([0.5, 0.5])),
+        [
+            dist.HalfNormal(1.0, validate_args=False),
+            dist.HalfNormal(2.0, validate_args=False),
+        ],
+        support=dist.constraints.real,
+        validate_args=False,
+    )
+    value = jnp.array([-1.0, 1.0])
+    assert jnp.isneginf(mixture.log_prob(value)[0])
+    grad = jax.grad(lambda x: mixture.log_prob(x).sum())(value)
+    assert not jnp.isnan(grad).any()