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: ***@***.***>

[esennesh]

Ok, plainly I need to re-fix the regex and assertions in the tests to handle the unwrapped case. I'll fold that into the second patch and force-push again and then we should have a minimal PR for which all the checks pass.

[github-actions]

Benchmark report

this PR bugfix/null_mixture_weights_grad at 2e891eda vs baseline master at 1f6877a0

  run time:     unchanged across 32 benchmarks
- compile time: 1 slower, 0 faster

Significant changes (1)

                                ──────── run time ───────     ────── compile time ──────
  benchmark                     baseline   this PR      Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────
- predictive_forward_sampling   725.9 ms  718.3 ms  -1.1%     160.9 ms  213.9 ms  +32.9%

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ─────── run time ───────     ────── compile time ─────
  benchmark                      baseline  this PR      Δ     baseline   this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            4.1 ms   4.0 ms  -1.6%     372.6 ms  377.0 ms  +1.2%
  categorical_log_prob             2.1 ms   2.2 ms  +4.8%      70.7 ms   68.9 ms  -2.6%
  dirichlet_log_prob               710 µs   675 µs  -5.0%     391.0 ms  386.1 ms  -1.2%
  dirichlet_sample                45.6 ms  45.9 ms  +0.6%     867.7 ms  849.1 ms  -2.1%
  gamma_log_prob                   2.1 ms   2.1 ms  -2.9%       2.07 s    2.05 s  -0.8%
  gamma_sample                    20.3 ms  20.6 ms  +1.1%     828.9 ms  834.6 ms  +0.7%
  lkj_cholesky_sample              5.4 ms   5.3 ms  -0.5%       1.18 s    1.21 s  +2.4%
  mixture_same_family_log_prob     2.1 ms   2.1 ms  +1.5%     107.6 ms  101.7 ms  -5.5%
  multivariate_normal_log_prob     277 µs   275 µs  -0.7%     168.0 ms  171.4 ms  +2.0%
  normal_log_prob                  556 µs   543 µs  -2.3%      59.0 ms   54.1 ms  -8.3%
  normal_sample                   20.7 ms  20.9 ms  +1.1%     199.1 ms  205.7 ms  +3.3%
  stick_breaking_transform         6.4 ms   6.2 ms  -2.3%     213.0 ms  205.9 ms  -3.3%
  student_t_log_prob               3.1 ms   3.1 ms  -0.7%      77.4 ms   80.8 ms  +4.4%
  truncated_normal_log_prob        677 µs   688 µs  +1.5%      53.7 ms   53.5 ms  -0.3%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    40.0 ms   40.9 ms     +2.0%       3.93 s    3.85 s   -2.0%
  log_density_hierarchical          3.6 ms    3.7 ms     +2.1%       1.25 s    1.23 s   -1.9%
  nested_handler_stack              1.4 ms    1.4 ms     +3.4%       421 µs    476 µs  +13.0%
  potential_energy_and_grad          21 µs     21 µs     -3.2%     102.7 ms   99.3 ms   -3.3%
- predictive_forward_sampling     725.9 ms  718.3 ms     -1.1%     160.9 ms  213.9 ms  +32.9%
  trace_seeded_model                818 µs    942 µs  (+15.2%)     577.7 ms  550.2 ms   -4.8%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    728.6 ms  735.0 ms  +0.9%       3.37 s   3.44 s  +2.1%
  nuts_dense_mass_funnel       1.17 s    1.17 s  -0.4%       2.63 s   2.75 s  +4.6%
  nuts_eight_schools           1.14 s    1.15 s  +0.4%       2.55 s   2.58 s  +1.0%
  nuts_hierarchical_glm        4.88 s    4.86 s  -0.5%       5.00 s   4.89 s  -2.2%
  nuts_logistic_regression     1.10 s    1.08 s  -2.3%       3.48 s   3.37 s  -3.1%
  nuts_vectorized_chains       2.53 s    2.51 s  -0.9%       2.95 s   2.88 s  -2.2%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 313.4 ms  310.7 ms  -0.9%       3.31 s   3.31 s  +0.1%
  svi_autodiagonalnormal_hierarchical          1.02 s    1.02 s  +0.4%       5.26 s   5.23 s  -0.6%
  svi_automultivariatenormal_eight_schools   736.8 ms  741.0 ms  +0.6%       4.05 s   4.15 s  +2.5%
  svi_autonormal_logistic                    760.9 ms  768.9 ms  +1.0%       3.65 s   3.57 s  -2.1%
  svi_multi_particle_elbo                      1.49 s    1.52 s  +1.7%       3.57 s   3.66 s  +2.5%
  svi_trace_mean_field_elbo                    1.30 s    1.29 s  -0.6%       5.52 s   5.55 s  +0.5%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	bugfix/null_mixture_weights_grad
commit	1f6877a0	2e891eda
numpyro	0.21.0	0.21.0
jax	0.11.0	0.11.0
backend	cpu	cpu
python	3.14.7	3.14.7

Runner: Linux-6.17.0-1020-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[esennesh]

I've corrected the patches on this branch and rebased atop master, but it looks like there's been a regression outside the code this branch touches :-/.

[esennesh]

Oh weird. Now the checks all pass. Neat! Two reviews requested from @juanitorduz and @Qazalbash , for what's now a much more minimal patch.