jax.debug_nans

jax.debug_nans = <jax._src.config.State object>

Context manager for jax_debug_nans config option.

Add nan checks to every operation. When a nan is detected on the output of a jit-compiled computation, call into the un-compiled version in an attempt to more precisely identify the operation which produced the nan.

Parameters:

new_val (Any)