NaN produced in code, absent in de-optimized function

[Piripuz]

Hello, I am quite new to all of this, so I'm sorry if I'm doing something wrong :)

I am writing a pretty complex MLE using quite intensively both JAX and JAXopt. Everything is working quite well, but when I try to find the gradient of the final likelihood (that is found itself by running some optimization and bisection algorithms) I find some NaN values.

Setting

jax.config.update("jax_enable_x64", True)
jax.config.update("jax_disable_jit", True)

the NaN values don't disappear, and setting

jax.config.update("jax_debug_nans", True)

I get the message

FloatingPointError: invalid value (nan) encountered in jit(convert_element_type). 
Because jax_config.debug_nans.value and/or config.jax_debug_infs is set,
the de-optimized function (i.e., the function as if the `jit` decorator were removed) was called in an attempt to get a more precise error message.
However, the de-optimized function did not produce invalid values during its execution. 
This behavior can result from `jit` optimizations causing the invalid value to be produced. 
It may also arise from having nan/inf constants as outputs, like `jax.jit(lambda ...: jax.numpy.nan)(...)`. 

It may be possible to avoid the invalid value by removing the `jit` decorator, at the cost of losing optimizations. 

If you see this error, consider opening a bug report at https://github.com/jax-ml/jax.

Note that this appears even after removing the jit decorator.

I am thus pretty blocked, and don't really even know where these nan values are coming from. What is a reasonable step to take, after all these "canonical" steps failed? Should I just refrain from the idea of managing to have some sort of gradient from such a complex function (that would be something pretty amazing to have, but clearly not fundamental)?

[dfm]

It's not that uncommon for numerical optimizers like JAXopt to end up in parts of parameter space that cause NaNs, so I'm not totally shocked by this outcome. You may need to tune the parameters of your optimizer to take smaller steps or reparameterize your problem to exclude invalid regions. It can be useful to try to identify a set of inputs where you can find NaNs without optimization to try to diagnose what's going on.

In the meantime, you could also try sharing the full stack trace from above that error message, but I expect tuning the algorithms will be the most fruitful!