Fix Beta with concentration1=1 gives nan log_prob at value=0 #2089
Conversation
juanitorduz
added
bug
numpyro/distributions/continuous.py
Outdated
| return self._dirichlet.log_prob(jnp.stack([value, 1.0 - value], -1)) | ||
| # Handle edge cases where concentration1=1 and value=0, or concentration0=1 and value=1 | ||
| # These cases would result in nan due to log(0) * 0 in the Dirichlet computation | ||
| log_prob = jnp.where( |
There was a problem hiding this comment.
interesting! could you check if grads w.r.t. value, concentration1, concentration0 are not NaN?
There was a problem hiding this comment.
Good catch 🙈 ! The gradients are NaN with this approach. After some investigation and talking with Claude Code 😅 . We arrived to a solution via https://docs.jax.dev/en/latest/_autosummary/jax.custom_jvp.html, see 13fc288 and 19f8be7
There was a problem hiding this comment.
It is better to mask the extreme values by canonical ones before calling direchlet log_prob. Using where is better than custom jvp I think: https://github.com/tensorflow/probability/blob/main/discussion/where-nan.pdf
There was a problem hiding this comment.
ok! I tried to use this approach in 1a9eae1
numpyro/distributions/continuous.py
Outdated
| # Step 4: Compute correct forward-pass value at boundaries | ||
| # Use stop_gradient to prevent gradients from flowing through this branch | ||
| # xlogy(0, 0) = 0 gives the correct value when concentration=1 at boundaries | ||
| boundary_log_prob = jax.lax.stop_gradient( |
There was a problem hiding this comment.
how about setting this to 0.0 instead of using stop_gradient?
There was a problem hiding this comment.
I could not make it work without stop_gradient :(
numpyro/distributions/continuous.py
Outdated
| safe_value = jnp.where(is_boundary, 0.5, value) | ||
|
|
||
| # Step 3: Compute log_prob with safe values (gradients flow through here) | ||
| safe_log_prob = ( |
There was a problem hiding this comment.
could we use self.dirichlet.log_prob here?
Closes #2088