code: def _backward_nan_bwd(
      residuals, grad_x, perturbed, x, name, terminate
  ):
      del residuals, perturbed
      msg = "   primals={x}\ncotangents={grad_x}"
      if name is not None:
          msg = f"{name}:\n" + msg
      jax.debug.print(  # pyright: ignore
          msg, x=_LongRepr(x), grad_x=_LongRepr(grad_x), ordered=True
      )
      if terminate:
          nans = [
              jnp.isnan(a).any()
              for a in jtu.tree_leaves(filter(grad_x, is_array_like))
          ]
          grad_x = error_if(
              grad_x, jnp.any(jnp.stack(nans)), "Encountered NaN"
          )
      return grad_x
  
file_path: equinox/equinox/debug/_backward_nan.py
imports: from typing import Any
  import jax
  import jax._src.traceback_util
  import jax.core
  import jax.numpy
  import jax.tree_util
  from .._ad import filter_custom_vjp
  from .._errors import error_if
  from .._filters import filter, is_array_like
  from .._module import Module
  from .._pretty_print import tree_pformat
label_desc: Use image node_id 4 for calling a global function with example usage: _backward_nan_bwd(residuals, grad_x, perturbed, x, name, terminate) and returns: grad_x