code: def checkpointed_while_loop(
      cond_fun: Callable[[_T], _Bool],
      body_fun: Callable[[_T], _T],
      init_val: _T,
      *,
      max_steps: Optional[int] = None,
      buffers: Optional[
          Callable[[_T], Union[_Node, Sequence[_Node]]]
      ] = None,
      checkpoints: Optional[int] = None,
  ) -> _T:
      """Reverse-mode autodifferentiable while loop, using optimal online checkpointing.
  
      The usual `jax.lax.while_loop` is not reverse-mode autodifferentiable, since it
      would need to save a potentially unbounded amount of residuals between the forward
      and backward pass. However, JAX/XLA requires that all memory buffers be of known
      (bounded) size.
  
      This works around this limitation by saving values to a prespecified number of
      checkpoints, and then recomputing other intermediate value on-the-fly.
  
      Checkpointing in this way is a classical autodifferentiation technique, usually used
      to reduce memory consumption. (And it's still useful for this purpose for us too.)
  
      **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. Set to `None` to allow an
          arbitrary number of steps. (`checkpointed_while_loop` is reverse-mode
          autodifferentiable regardless of whether `max_steps` is finite.)
      - `checkpoints`: The number of steps at which to checkpoint. The memory consumed
          will be that of `checkpoints`-many copies of `init_val`. (As the state is
          updated throughout the loop.)
      - `buffers`: If passed, then every array in `tree_leaves(buffers(init_val))` will
          become a write-only buffer. (Supporting only `.at[].set()`.)
  
      **Returns:**
  
      The final value; as `lax.while_loop`.
  
      !!! Info
  
          This function is not forward-mode autodifferentiable.
  
      !!! Info
  
          `buffers` is useful in the following way.
  
          Recall that a checkpointed while loop works by storing a copy of the evolving
          value (of the same structure as `init_val`). However, if part of that value is
          a large buffer than is being progressively written in to as the loop progresses,
          then we will end up with multiple copies of this large buffer.
  
          As this slows things down, we offer a special API specifically to handle this
          case.
  
      !!! Danger
  
          Note that `buffers` is subject to the following restrictions:
  
          - You should never write to the same location twice.
          - You should only read from it (`buf[i]`) at locations (`i`) that you have
            written to previously.
  
          These assumptions are *completely unchecked* and you will get incorrect
          gradients if you violate these assumptions.
  
      ??? cite "References"
  
          Selecting which steps at which to save checkpoints (and when this is done, which
          old checkpoint to evict) is important for minimising the amount of recomputation
          performed.
  
          This is a difficult, but solved, problem! So if you are using this function in
          academic work, then you should cite the following references.
  
          The implementation here performs "online checkpointing", as the number of steps
          is not known in advance. This was developed in:
  
          ```bibtex
          @article{stumm2010new,
              author = {Stumm, Philipp and Walther, Andrea},
              title = {New Algorithms for Optimal Online Checkpointing},
              journal = {SIAM Journal on Scientific Computing},
              volume = {32},
              number = {2},
              pages = {836--854},
              year = {2010},
              doi = {10.1137/080742439},
          }
  
          @article{wang2009minimal,
              author = {Wang, Qiqi and Moin, Parviz and Iaccarino, Gianluca},
              title = {Minimal Repetition Dynamic Checkpointing Algorithm for Unsteady
                       Adjoint Calculation},
              journal = {SIAM Journal on Scientific Computing},
              volume = {31},
              number = {4},
              pages = {2549--2567},
              year = {2009},
              doi = {10.1137/080727890},
          }
          ```
  
          For reference, the classical "offline checkpointing" (also known as "treeverse",
          "recursive binary checkpointing", "revolve" etc.) was developed in:
  
          ```bibtex
          @article{griewank1992achieving,
              author = {Griewank, Andreas},
              title = {Achieving logarithmic growth of temporal and spatial complexity in
                       reverse automatic differentiation},
              journal = {Optimization Methods and Software},
              volume = {1},
              number = {1},
              pages = {35--54},
              year  = {1992},
              publisher = {Taylor & Francis},
              doi = {10.1080/10556789208805505},
          }
  
          @article{griewank2000revolve,
              author = {Griewank, Andreas and Walther, Andrea},
              title = {Algorithm 799: Revolve: An Implementation of Checkpointing for the
                       Reverse or Adjoint Mode of Computational Differentiation},
              year = {2000},
              publisher = {Association for Computing Machinery},
              volume = {26},
              number = {1},
              doi = {10.1145/347837.347846},
              journal = {ACM Trans. Math. Softw.},
              pages = {19--45},
          }
          ```
      """
      if checkpoints is None:
          if max_steps is None:
              raise ValueError(
                  "Must specify either `max_steps` or `checkpoints` in "
                  "`equinox.internal.checkpointed_while_loop`."
              )
          # Binomial logarithmic growth is what is needed in classical treeverse.
          #
          # If
          # `max_steps < (checkpoints + 1)(checkpoints + 2)/2`
          # then the time spend recomputing will be optimised; see equation (2.2) of
          # "New Algorithms for Optimal Online Checkpointing", Stumm and Walther 2010.
          # https://tu-dresden.de/mn/math/wir/ressourcen/dateien/forschung/publikationen/pdf2010/new_algorithms_for_optimal_online_checkpointing.pdf
          # and also the `_should_save_residuals` function below.
          #
          # So by default we use `checkpoints = O(sqrt(max_steps))`.
          # (Classical treeverse uses only `O(log(max_steps))`, of course, but doesn't
          # handle the online case.)
          if max_steps == 1:
              checkpoints = 1
          else:
              checkpoints = (
                  math.floor(-1.5 + math.sqrt(2 * max_steps + 0.25)) + 1
              )
              assert 2 * max_steps < (
                  (checkpoints + 1) * (checkpoints + 2)
              )
      if checkpoints < 1:
          raise ValueError("Must have at least one checkpoint")
      init_val = jtu.tree_map(fixed_asarray, init_val)
      if max_steps == 0:
          return init_val
      if max_steps is not None:
          checkpoints = min(checkpoints, max_steps)
      cond_fun_, body_fun_, init_val_, buffers_ = common_rewrite(
          cond_fun,
          body_fun,
          init_val,
          max_steps,
          buffers,
          makes_false_steps=False,
      )
      del cond_fun, body_fun, init_val, buffers
      cond_fun_ = filter_closure_convert(cond_fun_, init_val_)
      cond_fun_ = jtu.tree_map(_stop_gradient, cond_fun_)
      body_fun_ = filter_closure_convert(body_fun_, init_val_)
      vjp_arg = (init_val_, body_fun_)
      _, _, _, final_val = _checkpointed_while_loop(
          vjp_arg, cond_fun_, checkpoints, buffers_, max_steps
      )
      return final_val
  
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 1 for calling a global function with example usage: checkpointed_while_loop(cond_fun, body_fun, init_val) and returns: final_val, init_val