code: def _load_from_checkpoint(
      step_grad_val,
      index,
      residual_steps,
      residuals,
      checkpoints,
      max_steps,
  ):
      """Loads a residual from the store of checkpoints."""
      # step_grad_val is the current location of grad_val.
      # index is the next currently empty slot for saving a residual.
      step_grad_val, index, residual_steps = nonbatchable(
          (step_grad_val, index, residual_steps)
      )
  
      # Subtract one to get the most recent residual, and then load it.
      # (Clip to zero just to not error from index == 0 on the very last step; the result
      # is unused in this case.)
      read_index = jnp.maximum(index - 1, 0)
      step_val2, val2 = jtu.tree_map(
          ft.partial(_scalar_index, read_index),
          (residual_steps, residuals),
      )
  
      # We may need to keep this residual around, and jump back to it multiple times.
      # (In which case index2 == index.) Or this may be the last time and we don't need to
      # save it any more. (In which case index2 == index - 1.)
      # If `step_val2 + 1 == step_grad_val2` then we're about to make a U-turn on the next
      # step, so we won't need to load from this checkpoint. (And as
      # `_load_from_checkpoint` is itself used within a U-turn, then in practice this
      # triggers whenever we get >1 U-turns back-to-back.)
      if checkpoints == max_steps:
          # Known to be symbolically true in this case; this is a slight performance
          # optimisation.
          pred = True
      else:
          pred = step_val2 + 1 == step_grad_val
      index2 = jnp.where(pred, read_index, index)
      step_val2, index2 = nonbatchable((step_val2, index2))
      return step_val2, val2, index2
  
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 16 for calling a global function with example usage: _load_from_checkpoint(step_grad_val, index, residual_steps, residuals, checkpoints, max_steps) and returns: step_val2,  val2,  index2