code: def branched_error_if(
      x: PyTree,
      pred: Bool[ArrayLike, "..."],
      index: Int[ArrayLike, "..."],
      msgs: Sequence[str],
      *,
      on_error: Literal[
          "default", "raise", "breakpoint", "nan"
      ] = "default",
  ) -> PyTree:
      """As [`equinox.error_if`][], but will raise one of
      several `msgs` depending on the value of `index`. If `index` is vmap'd, then the
      error message from the largest value (across the whole batch) will be used.
      """
      leaves = jtu.tree_leaves((x, pred, index))
      # This carefully does not perform any JAX operations if `pred` and `index` are
      # a bool and an int.
      # This ensures we can use `error_if` before init_google.
      if any(is_array(leaf) for leaf in leaves):
          return branched_error_if_impl_jit(
              x, pred, index, msgs, on_error=on_error
          )
      else:
          return branched_error_if_impl(
              x, pred, index, msgs, 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 5 for calling a global function with example usage: branched_error_if(x, pred, index, msgs) and returns: branched_error_if_impl_jit, branched_error_if_impl