code: def branched_error_if_impl(
      x: PyTree,
      pred: Bool[ArrayLike, "..."],
      index: Int[ArrayLike, "..."],
      msgs: Sequence[str],
      *,
      on_error: Literal["default", "raise", "breakpoint", "nan"],
  ) -> PyTree:
      if on_error == "default":
          on_error = EQX_ON_ERROR
      elif on_error not in ("raise", "breakpoint", "nan"):
          raise RuntimeError("Unrecognised value for `on_error`.")
      with jax.ensure_compile_time_eval():
          # 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 not isinstance(pred, bool):
              pred = unvmap_any(pred)
          if not isinstance(index, int):
              index = unvmap_max(index)
          if not isinstance(pred, jax.core.Tracer):
              if isinstance(pred, Array):
                  pred = pred.item()
              assert type(pred) is bool
              if pred:
                  if not isinstance(index, jax.core.Tracer):
                      if isinstance(index, Array):
                          index = index.item()
                      assert type(index) is int
                      if on_error == "raise":
                          raise EquinoxTracetimeError(msgs[index])
                      elif on_error == "breakpoint":
                          print(msgs[index])
                          breakpoint()
                      elif on_error == "nan":
                          warnings.warn(
                              "Resolving error at trace time (because the predicate is "
                              "statically resolvable), by substituting NaNs (because "
                              "`on_error='nan'`)."
                          )
                          return jtu.tree_map(_nan_like, x)
                      else:
                          assert False
                  # else defer error to runtime, when the index is known.
              else:
                  return x
  
      dynamic_x, static_x = partition(x, is_array)
      flat = jtu.tree_leaves(dynamic_x)
      if len(flat) == 0:
          raise ValueError("No arrays to thread error on to.")
      dynamic_x = _error(
          dynamic_x, pred, index, msgs=msgs, on_error=on_error
      )
      return combine(dynamic_x, static_x)
  
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 6 for calling a global function with example usage: branched_error_if_impl(x, pred, index, msgs) and returns: combine, x, jtu