Inconsistent nan behavior across devices

[cagrikymk]

Hi,
I have a fairly complex code that calculates d(dy/dx)/dp where y is a scalar output, x is the input and p are the parameters to optimize. While everything works perfectly on 1080TI and V100 gpus, I get NaN values for the gradients on A100 GPUs. I am using double precision.

I wonder what the best way is to debug such a problem. I assume this is probably related to architecture specific XLA optimizations.

[patrick-kidger]

equinox.internal.debug_backward_nan (https://github.com/patrick-kidger/equinox/tree/main/equinox/internal) is probably a good place to start.

You might also try setting the environment variable JAX_DEBUG_NANS=1.