brainstate.transform.debug_nan_if

Contents

brainstate.transform.debug_nan_if#

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

Conditionally run fn with NaN / Inf detection.

Equivalent to:

if has_nan:
    debug_nan(fn, *args, phase=phase)

but JIT-compatible via jax.lax.cond.

Parameters:

  • has_nan (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Condition to trigger debugging. Batched arrays are collapsed with unvmap(..., op='any').

  • fn (Callable) – The function to debug.

  • *args – Arguments to pass to the function.

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

Return type:

None

See also

debug_nan

Unconditional variant.

Contents