code: def error_if(
      x: PyTree,
      pred: Bool[ArrayLike, "..."],
      msg: str,
      *,
      on_error: Literal[
          "default", "raise", "breakpoint", "nan"
      ] = "default",
  ) -> PyTree:
      """Throws an error based on runtime values. Works even under JIT.
  
      **Arguments:**
  
      - `x`: will be returned unchanged. This is used to determine where the error check
          happens in the overall computation: it will happen after `x` is computed and
          before the return value is used. `x` can be any PyTree, and it must contain at
          least one array.
      - `pred`: a boolean for whether to raise an error. Can be an array of bools; an
          error will be raised if any of them are `True`. If vmap'd then an error will be
          raised if any batch element has `True`.
      - `msg`: the string to display as an error message.
  
      In addition, the `EQX_ON_ERROR` environment variable is checked for how any runtime
      errors should be handled. Possible values are:
  
      - `EQX_ON_ERROR=raise` will raise a runtime error.
      - `EQX_ON_ERROR=nan` will return `NaN` instead of `x`, and then continue the
          computation.
      - `EQX_ON_ERROR=breakpoint` will open a debugger.
          - Note that this option may prevent certain compiler optimisations, so
              permanently fixing this value is not recommended.
          - You will need to also pass the `-s` flag to `pytest`, if you are
              also using that.
          - This will sometimes raise a trace-time error due to JAX bug
              [#16732](https://github.com/google/jax/issues/16732). (Bugs whilst debugging
              bugs, eek!) If this happens, then it can be worked around by additionally
              setting the `EQX_ON_ERROR_BREAKPOINT_FRAMES` variable to a small integer,
              which specifies how many frames upwards the debugger should capture. The
              JAX bug is triggered when taking too many frames.
  
      After changing an environment variable, the Python process must be restarted.
  
      **Returns:**
  
      The original argument `x` unchanged. **If this return value is unused then the error
      check will not be performed.** (It will be removed as part of dead code
      elimination.)
  
      !!! Example
  
          ```python
          @jax.jit
          def f(x):
              x = error_if(x, x < 0, "x must be >= 0")
              # ...use x in your computation...
              return x
  
          f(jax.numpy.array(-1))
          ```
      """
      return branched_error_if(x, pred, 0, [msg], on_error=on_error)
  
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 4 for calling a global function with example usage: error_if(x, pred, msg) and returns: branched_error_if