code: def breakpoint_if(pred: Bool[Array, "..."], **kwargs):
      """As `jax.debug.breakpoint`, but only triggers if `pred` is True.
  
      **Arguments:**
  
      - `pred`: the predicate for whether to trigger the breakpoint.
      - `**kwargs`: any other keyword arguments to forward to `jax.debug.breakpoint`.
  
      **Returns:**
  
      Nothing.
      """
      # We can't just write `jax.debug.breakpoint` for the second branch. For some reason
      # it needs as lambda wrapper.
      token = kwargs.get("token", None)
      return lax.cond(
          unvmap_any(pred),
          lambda: jax.debug.breakpoint(**kwargs),
          lambda: token,
      )
  
imports: import jax
  import jax._src.traceback_util
  import jax.core
  import jax.lax
  from jaxtyping import Array, Bool
  from .._unvmap import unvmap_any
label_desc: Use image node_id 1 for calling a global function with example usage: breakpoint_if(pred) and returns: lax