code: def _while_loop(cond_fun, body_fun, val, max_steps, base):
      if max_steps == 1:
          return body_fun(val)
      else:
  
          def call(val):
              return _while_loop(
                  cond_fun, body_fun, val, max_steps // base, base
              )
  
          def scan_fn(val, _):
              return (
                  lax.cond(cond_fun(val), call, lambda x: x, val),
                  None,
              )
  
          # Don't put checkpointing on the lowest level
          if max_steps != base:
              scan_fn = jax.checkpoint(
                  scan_fn, prevent_cse=False
              )  # pyright: ignore
  
          return lax.scan(scan_fn, val, xs=None, length=base)[0]
  
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 2 for calling a global function with example usage: _while_loop(cond_fun, body_fun, val, max_steps, base) and returns: body_fun, lax, _while_loop, lax,  None