code: def test_backward_nan(capfd):
      @eqx.filter_custom_vjp
      def backward_nan(x):
          return x
  
      @backward_nan.def_fwd
      def backward_nan_fwd(perturbed, x):
          del perturbed
          return backward_nan(x), None
  
      @backward_nan.def_bwd
      def backward_nan_bwd(residual, grad_x, perturbed, x):
          del residual, grad_x, perturbed, x
          return jnp.nan
  
      @eqx.filter_jit
      @jax.grad
      def f(x, terminate):
          y = eqx.debug.backward_nan(x, name="foo", terminate=terminate)
          return backward_nan(y)
  
      capfd.readouterr()
      f(jnp.array(1.0), terminate=False)
      jax.effects_barrier()
      text, _ = capfd.readouterr()
      assert (
          text
          == "foo:\n   primals=array(1., dtype=float32)\ncotangents=array(nan, dtype=float32)\n"  # noqa: E501
      )
  
      with pytest.raises(Exception):
          f(jnp.array(1.0), terminate=True)
  
file_path: equinox/tests/test_debug.py
imports: import equinox
  import jax
  import jax.numpy
  import jax.tree_util
  import pytest
label_desc: Use image node_id 1 for calling a global function with example usage: test_backward_nan(capfd) and returns: x, backward_nan,  None, jnp, backward_nan