fix NaN mixture log_prob gradient at zero weights

[esennesh]

_MixtureBase.log_prob wrote the density in log-weight form via log_softmax(mixing.logits) + log p_k. For a probs-parameterized mixing Categorical with an exact-zero weight, logits takes log(0): the forward value was masked to finite but the VJP evaluated 1/0 = inf, producing a NaN gradient into every parameter feeding the weights.

[Qazalbash]

Hi @esennesh, thank you. Do you have an MRE where this bug was appearing?

[esennesh]

Hi @esennesh, thank you. Do you have an MRE where this bug was appearing?

The below should trigger it. That's minimal, of course. My actual use-case was a Gaussian mixture model in which the means and weights were dynamically determined by upstream latent variables, so that sometimes zero weights would occur in some components and sometimes not. Since we're working in Jax, the dynamic control flow to just leave out all zero-weight components is a bit finicky, if it's possible.

import jax.numpy as jnp
import numpyro.distributions as dist

def log_prob(probs):
    mixing = dist.Categorical(probs=probs)
    components = dist.Normal(jnp.array([0.0, 1.0, 2.0]), 1.0)
    return dist.MixtureSameFamily(mixing, components).log_prob(0.3)

probs = jnp.array([0.0, 0.5, 0.5])  # one component has exactly zero weight

print(log_prob(probs))            # finite either way: ~ -1.359
print(jax.grad(log_prob)(probs))  # master: [nan nan nan] — fixed: [1.454, 0.960, 0.577]```

[Qazalbash]

@esennesh, thanks for the MRE. I have modified it a little.

import jax
import jax.numpy as jnp

import numpyro.distributions as dist


def log_prob(probs):
    mixing_distribution = dist.Categorical(probs=probs, validate_args=True)
    component_distributions = dist.Normal(
        loc=jnp.array([0.0, 1.0, 2.0]),
        scale=1.0,
        validate_args=True,
    )
    return dist.MixtureSameFamily(
        mixing_distribution,
        component_distributions,
        validate_args=True,
    ).log_prob(0.3)


log_prob_jit = jax.jit(log_prob)

mixing_probs = jnp.array([0.0, 0.5, 0.5])

log_prob_val = log_prob(mixing_probs)
print(log_prob_val)

log_prob_val = log_prob_jit(mixing_probs)
print(log_prob_val)

with jax.debug_nans(True):
    grad_log_prob_val = jax.grad(log_prob)(mixing_probs)
    print(grad_log_prob_val)

with jax.debug_nans(True):
    grad_log_prob_val = jax.grad(log_prob_jit)(mixing_probs)
    print(grad_log_prob_val)

It is reproducing the same error. This is the output before fix:

-1.5938032
-1.5938032
Traceback (most recent call last):
  File "/home/gradf/academia/numpyro/test.py", line 32, in <module>
    grad_log_prob_val = jax.grad(log_prob)(mixing_probs)
  File "/home/gradf/academia/numpyro/test.py", line 18, in log_prob
    ).log_prob(0.3)
  File "/home/gradf/academia/numpyro/numpyro/distributions/util.py", line 799, in wrapper
    log_prob = log_prob_fn(self, *args, **kwargs)
  File "/home/gradf/academia/numpyro/numpyro/distributions/mixtures.py", line 170, in log_prob
    sum_log_probs = self.component_log_probs(value)
  File "/home/gradf/academia/numpyro/numpyro/distributions/mixtures.py", line 296, in component_log_probs
    return jax.nn.log_softmax(self.mixing_distribution.logits) + component_log_probs
  File "/home/gradf/academia/numpyro/numpyro/distributions/util.py", line 790, in __get__
    value = self.wrapped(instance)
  File "/home/gradf/academia/numpyro/numpyro/distributions/discrete.py", line 722, in logits
    return _to_logits_multinom(self.probs)
  File "/home/gradf/academia/numpyro/numpyro/distributions/discrete.py", line 72, in _to_logits_multinom
    return jnp.clip(jnp.log(probs), minval)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: FloatingPointError: invalid value (nan) encountered in div

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/gradf/academia/numpyro/test.py", line 32, in <module>
    grad_log_prob_val = jax.grad(log_prob)(mixing_probs)
FloatingPointError: invalid value (nan) encountered in div
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

and after fix:

-1.5938032
-1.5938032
[ 0.          0.53704965 -0.53704953]
[ 0.          0.53704965 -0.53704953]

I assume the numerical answers are correct. And the fix I have proposed is,

def _to_logits_multinom(probs: ArrayLike) -> ArrayLike:
-    minval = jnp.finfo(jnp.result_type(probs)).min
-    return jnp.clip(jnp.log(probs), minval)
+    safe_probs = jnp.where(probs > 0, probs, 1.0)
+    safe_log_probs = jnp.where(probs > 0, jnp.log(safe_probs), -jnp.inf)
+    return safe_log_probs

@juanitorduz what are your thoughts?

[fehiepsi]

Do I understand correctly that the issue can be addressed with simpler change?

[Qazalbash]

Do I understand correctly that the issue can be addressed with simpler change?

Yes. The MRE does pass with the little change. Although I have not tested it on the rest of the test suite.

[Qazalbash]

Hi @esennesh, have you had a chance to try my fix in your workflow to see if it is working?

[esennesh]

Doing that tonight or tomorrow morning, most likely.

…

On Sat, Aug 1, 2026, 6:52 AM Meesum Qazalbash ***@***.***> wrote: *Qazalbash* left a comment (pyro-ppl/numpyro#2224) <#2224 (comment)> Hi @esennesh <https://github.com/esennesh>, have you had a chance to try my fix in your workflow to see if it is working? — Reply to this email directly, view it on GitHub <#2224?email_source=notifications&email_token=AAFRWFNCDBPQL2S5WL4ZVA35HXYYJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJVGE3TCMJQGU32M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5151711057>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFRWFMD6J6HP4GT5LANXHD5HXYYJAVCNFSNUABFKJSXA33TNF2G64TZHMYTOMBVHAYDKNBQHNEXG43VMU5TIOJTGQZDMNBTHAZ2C5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAFRWFMDOTK5OR4AUMEHEFD5HXYYJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJVGE3TCMJQGU32M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAFRWFPCHQPKH4RWXICNZAD5HXYYJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJVGE3TCMJQGU32M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>