Grads w.r.t. weights of MixtureGeneral Distribution are giving nans

[Qazalbash]

Hi,

We have created some models where we estimate the weights of the MixtureGeneral distribution. However, when computing the gradient of this argument, we are encountering nan values. We enabled jax.config.update("debug_nan", True) to diagnose the issue, and it pointed to the following line:

numpyro/numpyro/distributions/mixtures.py

Line 152 in 8e9313f

return jax.nn.logsumexp(sum_log_probs, axis=-1)

I suspect that after the implementation of #1791, extra care is needed to handle inf and nan values, possibly by using a double where for a safe logsumexp.

Important

This is an urgent issue, so a prompt response would be greatly appreciated.

[fehiepsi]

You can add jax.debug.print(...) to inspect the component log probs. If all of the component log probs are -inf, nan will happen.

[Qazalbash]

I am able to get gradients of weights even with -jnp.inf, by modifying,

numpyro/numpyro/distributions/mixtures.py

Lines 148 to 152 in 8e9313f

	@validate_sample
	def log_prob(self, value, intermediates=None):
	del intermediates
	sum_log_probs = self.component_log_probs(value)
	return jax.nn.logsumexp(sum_log_probs, axis=-1)

to

@validate_sample
def log_prob(self, value, intermediates=None):
    del intermediates
    sum_log_probs = self.component_log_probs(value)
    safe_sum_log_probs = jnp.where(
        jnp.isneginf(sum_log_probs), -jnp.inf, sum_log_probs
    )
    return jax.nn.logsumexp(safe_sum_log_probs, axis=-1)