brainstate.transform.debug_nan

Contents

brainstate.transform.debug_nan#

brainstate.transform.debug_nan(fn, *args, phase='')[source]#

Run fn with NaN / Inf detection (JIT-compatible).

An error is raised only when NaN/Inf reaches fn’s observable outputs, so NaN that is computed but masked away does not trigger a false alarm.

Parameters:

  • fn (Callable) – The function to debug.

  • *args – Arguments to pass to the function.

  • phase (str) – Label prepended to the error message.

Return type:

None

Raises:

RuntimeError – If NaN or Inf contaminates fn’s outputs or updated state.

See also

debug_nan_if

Conditional variant.

Examples

>>> import jax.numpy as jnp
>>> import brainstate
>>> brainstate.transform.debug_nan(lambda x: jnp.log(x), jnp.array([-1.0, 1.0]))
Traceback (most recent call last):
    ...
RuntimeError: NaN/Inf detected ...
Contents