code: def _error(x, pred, index, *, msgs, on_error):
      if on_error == "raise":
  
          def raises(_index):
              raise EqxRuntimeError(msgs[_index.item()])
  
          def tpu_msg(_out, _index):
              msg = msgs[_index.item()]
              # `print` doesn't work; nor does `jax.debug.print`.
              # But both `input` and `jax.debug.breakpoint` do. The former allows us to
              # actually display something to the user.
              input(msg + _tpu_msg)
              # We do the tree_map inside the pure_callback, not outside, so that `out`
              # has a data dependency and doesn't get optimised out.
              return jtu.tree_map(_nan_like, _out)
  
          def handle_error():  # pyright: ignore
              out = jax.pure_callback(
                  raises, struct, index
              )  # pyright: ignore
              # If we make it this far then we're on the TPU, which squelches runtime
              # errors and returns dummy values instead.
              # Fortunately, we're able to outsmart it!
              return jax.pure_callback(
                  tpu_msg, struct, out, index
              )  # pyright: ignore
  
          struct = jax.eval_shape(lambda: x)
          return lax.cond(pred, handle_error, lambda: x)
  
      elif on_error == "breakpoint":
  
          def display_msg(_index):
              print(msgs[_index.item()])
              return _index
  
          def to_nan(_index):
              del _index
              return jtu.tree_map(_nan_like, struct)
  
          def handle_error():
              index_struct = jax.eval_shape(lambda: index)
              _index = jax.pure_callback(  # pyright: ignore
                  display_msg, index_struct, index, vectorized=True
              )
              # Support JAX with and without DCE behaviour on breakpoints.
              breakpoint_params = inspect.signature(
                  jax.debug.breakpoint
              ).parameters.keys()
              breakpoint_kwargs = {}
              if "token" in breakpoint_params:
                  breakpoint_kwargs["token"] = _index
              if "vectorized" in breakpoint_params:
                  breakpoint_kwargs["vectorized"] = True
              if EQX_ON_ERROR_BREAKPOINT_FRAMES is not None:
                  breakpoint_kwargs[
                      "num_frames"
                  ] = EQX_ON_ERROR_BREAKPOINT_FRAMES
              _index = jax.debug.breakpoint(**breakpoint_kwargs)
              return jax.pure_callback(  # pyright: ignore
                  to_nan, struct, _index, vectorized=True
              )
  
          struct = jax.eval_shape(lambda: x)
          return lax.cond(pred, handle_error, lambda: x)
  
      elif on_error == "nan":
          return lax.cond(
              pred, ft.partial(jtu.tree_map, _nan_like), lambda y: y, x
          )
      else:
          assert False
  
imports: import functools
  import inspect
  import warnings
  from collections.abc import Sequence
  from typing import Literal, Union
  import jax
  import jax._src.traceback_util
  import jax.core
  import jax.lax
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, ArrayLike, Bool, Int, PyTree
  from ._ad import filter_custom_jvp
  from ._config import EQX_ON_ERROR, EQX_ON_ERROR_BREAKPOINT_FRAMES
  from ._doc_utils import doc_remove_args
  from ._filters import combine, is_array, partition
  from ._jit import filter_jit
  from ._misc import currently_jitting
  from ._unvmap import unvmap_any, unvmap_max
label_desc: Use image node_id 2 for calling a global function with example usage: _error(x, pred, index) and returns: lax, jtu, jax, lax, _index, jtu, jax, lax