code: def bounded_while_loop(
      cond_fun: Callable[[_T], _Bool],
      body_fun: Callable[[_T], _T],
      init_val: _T,
      *,
      max_steps: int,
      buffers: Optional[
          Callable[[_T], Union[_Node, Sequence[_Node]]]
      ] = None,
      base: int = 16,
  ):
      """Reverse-mode autodifferentiable while loop.
  
      **Arguments:**
  
      - `cond_fun`: As `lax.while_loop`.
      - `body_fun`: As `lax.while_loop`.
      - `init_val`: As `lax.while_loop`.
      - `max_steps`: A bound on the maximum number of steps, after which the loop
          terminates unconditionally. Can be set to `None` for arbitrarily many steps.
      - `buffers`: If passed, then every leaf of `tree_leaves(buffers(init_val))` must
          be an array; all such arrays become buffers supporting only `[]` and
          `.at[].set()`. However they will act efficiently, without spurious copies.
      - `base`: Run time will increase slightly as `base` increases. Compilation time will
          decrease substantially as `math.ceil(math.log(max_steps, base))` decreases.
          (Which happens as `base` increases.)
  
      **Returns:**
  
      The final value; as `lax.while_loop`.
  
      """
  
      if not isinstance(max_steps, int) or max_steps < 0:
          raise ValueError("max_steps must be a non-negative integer")
      init_val = jtu.tree_map(fixed_asarray, init_val)
      if max_steps == 0:
          return init_val
  
      cond_fun_, body_fun_, init_val_, _ = common_rewrite(
          cond_fun,
          body_fun,
          init_val,
          max_steps,
          buffers,
          makes_false_steps=True,
      )
      del cond_fun, body_fun, init_val
      rounded_max_steps = base ** int(
          math.ceil(math.log(max_steps, base))
      )
      _, _, _, val = _while_loop(
          cond_fun_, body_fun_, init_val_, rounded_max_steps, base
      )
      return val
  
file_path: equinox/equinox/internal/_loop/bounded.py
imports: import math
  from collections.abc import Callable, Sequence
  from typing import Any, Optional, TypeVar, Union
  import jax
  import jax.lax
  import jax.tree_util
  from jaxtyping import Array, Bool
  from .common import common_rewrite, fixed_asarray
label_desc: Use image node_id 1 for calling a global function with example usage: bounded_while_loop(cond_fun, body_fun, init_val) and returns: val, init_val