Closed Theo-Cheynel closed 1 year ago
Patch and project coverage have no change.
Comparison is base (
6903b24
) 100.00% compared to head (c473fc0
) 100.00%.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Do you have feedback about the report comment? Let us know in this issue.
Thanks!
Why ?
In
SO3.exp
, there is a shim to prevent NaNs by using a jnp.where with a Taylor approximation for values under epsilon. However, this is still catched by the jax_debug_nans configuration, even when there are no NaNs present in the final outputs (the NaNs appear in the intermediate computations).Contents of this PR
This can be solved easily by changing the arbitrary value from 0.0 to 1.0, thus removing the divisions by 0 and preventing the catch by the JAX NaN debugger.