code: def _checkpointed_while_loop_fwd(
      perturbed, vjp_arg, cond_fun, checkpoints, buffers, max_steps
  ):
      """Run the while loop, saving checkpoints whenever the controller
      (`_should_save_residual`) requires.
      """
      del perturbed
      init_val, body_fun = vjp_arg
      # Equation (2.2) of Stumm and Walther
      u2_minus_1 = ((checkpoints + 1) * (checkpoints + 2)) // 2 - 1
  
      def _cond_fun(carry):
          _, _, val, _, _ = carry
          return cond_fun(val)
  
      def _body_fun(carry):
          step, save_state, val, residual_steps, residuals = carry
          save_state, residual_steps = nonbatchable(
              (save_state, residual_steps)
          )
  
          step2 = step + 1
          save_residual, index, save_state2 = _should_save_residual(
              step,
              save_state,
              residual_steps,
              u2_minus_1,
              checkpoints,
              max_steps,
          )
          val2 = body_fun(val)
  
          def _maybe_update(xs, x):
              where_x = jnp.where(
                  save_residual, x, _scalar_index(index, xs)
              )
              where_x = cast(Array, where_x)
              return lax.dynamic_update_index_in_dim(
                  xs, where_x, index, axis=0
              )
  
          val_no_buffers = tree_at(
              buffers(None), val, replace_fn=_array_to_none
          )
          residual_steps2, residuals2 = jtu.tree_map(
              _maybe_update,
              (residual_steps, residuals),
              (step, val_no_buffers),
          )
          save_state2, residual_steps2 = nonbatchable(
              (save_state2, residual_steps2)
          )
          return step2, save_state2, val2, residual_steps2, residuals2
  
      int_dtype = (
          jnp.int64 if jax.config.jax_enable_x64 else jnp.int32
      )  # pyright: ignore
      init_step = jnp.array(
          0, dtype=int_dtype
      )  # dtype matches init_residual_steps
      init_save_state_sw_i = 0, checkpoints, checkpoints, True
      dtype_max = jnp.iinfo(int_dtype).max  # pyright: ignore
      init_save_state_wm = (
          jnp.zeros(checkpoints, dtype=int_dtype).at[0].set(dtype_max),
          jnp.full((checkpoints,), False),
      )
      init_save_state = (init_save_state_sw_i, init_save_state_wm)
      # Uses the fact that `_unreachable_checkpoint_step` returns intmax, so that in our
      # sorting later, in the steps < checkpoints case, all unused memory gets sorted
      # to the end.
      init_residual_steps = jnp.full(
          checkpoints,
          _unreachable_checkpoint_step(int_dtype),
          dtype=int_dtype,
      )
  
      init_val_no_buffers = tree_at(
          buffers(None), init_val, replace_fn=_array_to_none
      )
      # Fill value for the memory isn't important.
      init_residuals = jtu.tree_map(
          lambda x: jnp.zeros((checkpoints,) + x.shape, x.dtype),
          init_val_no_buffers,
      )
      init_carry = (
          init_step,
          init_save_state,
          init_val,
          init_residual_steps,
          init_residuals,
      )
      # `step` is an increasing counter 0, 1, 2, 3, ...
      # `save_state` is state used by the logic for whether to save a checkpoint each step
      # `val` is the evolving state of the loop
      # `residual_steps` is the buffer of the `step` for each checkpoint
      # `residuals` is the buffer of the `val` for each checkpoint
  
      while_loop = jax.named_call(
          lax.while_loop, name="checkpointed-fwd"
      )
      final_carry = while_loop(_cond_fun, _body_fun, init_carry)
      (
          num_steps,
          _,
          final_val,
          final_residual_steps,
          final_residuals,
      ) = final_carry
      filled_buffers = buffers(_is_none)(final_val)
  
      # The above procedure may produce residuals saved in jumbled-up order. Meanwhile
      # treeverse (used on the backward pass) treats the residuals like a stack,
      # reading and writing the most recent residual to and from the end. So sort the
      # residuals we've produced here to obtain the desired invariant, i.e. that the
      # residuals are in order.
      if checkpoints != max_steps:
          # If `checkpoints == max_steps` then residuals are already sorted.
          sort_indices = jnp.argsort(final_residual_steps)
          final_residual_steps, final_residuals = jtu.tree_map(
              ft.partial(_unique_index, sort_indices),
              (final_residual_steps, final_residuals),
          )
      num_steps, final_residual_steps = nonbatchable(
          (num_steps, final_residual_steps)
      )
      return final_val, (
          num_steps,
          final_residual_steps,
          final_residuals,
          filled_buffers,
      )
  
imports: import functools
  import math
  import operator
  import typing
  from collections.abc import Callable, Sequence
  from typing import Any, cast, Optional, TypeVar, Union
  import jax
  import jax.core
  import jax.lax
  import jax.numpy
  import jax.tree_util
  from jaxtyping import Array, ArrayLike, Bool
  from ..._ad import filter_closure_convert, filter_custom_vjp
  from ..._errors import error_if
  from ..._filters import combine, filter, is_array, is_inexact_array, partition
  from ..._module import Static
  from ..._tree import tree_at, tree_equal
  from .._nontraceable import nonbatchable
  from .common import common_rewrite, fixed_asarray
label_desc: Use image node_id 15 for calling a global function with example usage: _checkpointed_while_loop_fwd(perturbed, vjp_arg, cond_fun, checkpoints, buffers, max_steps) and returns: final_val,  num_steps,  final_residual_steps,  final_residuals,  filled_buffers, cond_fun, step2,  save_state2,  val2,  residual_steps2,  residuals2, lax