code: def _calc_next_checkpoint(
      step_val, step_grad_val, index, checkpoints
  ):
      """Determines the step at which we next want to save a checkpoint."""
      # Note that when this function is called, `step_val` is always at the most recent
      # checkpoint.
      step_val, step_grad_val, index = nonbatchable(
          (step_val, step_grad_val, index)
      )
  
      # Using treeverse...
      # ...Checkpoints are either placed binomially (most of the time)...
      out_binomial = step_val + (step_grad_val - step_val) // 2
      # ...or linearly (when the space to cross fits within the checkpoint budget).
      out_linear = step_val + 1
      within_budget = (step_grad_val - step_val - 2) <= (
          checkpoints - index
      )
      out = jnp.where(within_budget, out_linear, out_binomial)
      # Why -2?
      # If `step_val + 1 == step_grad_val` then we're just going to make a single U-turn,
      # and don't need to store any checkpoints.
      # If `step_val + 2 == step_grad_val` then we're just going make a fwd, then a
      # U-turn, and can then load the existing checkpoint at `step_val` in order to make
      # the next U-turn. Once again we don't need to store any checkpoints.
      # If `step_val + 3 == step_grad_val` then ideally we would store a checkpoint at
      # `step_val + 1`.
      # Meanwhile, `checkpoints - index` is the number of spaces we have left in which to
      # store checkpoints. (e.g. `checkpoints == index` indicates that our buffer is
      # full.)
      # Thus if `step_val + 3 == step_grad_val` then
      # `step_grad_val - step_val - 2 <= checkpoints - index`
      # is the desired condition. (For "save a checkpoint if you can".)
      # Now proceed by induction: if `step_val + i == step_grad_val` for some `i > 3` and
      # we trigger this condition, we will reduce the LHS by one (as `step_val `increments
      # in a fwd step) and increase the RHS side by one (as `index` increments as we save
      # a checkpoint) and we eventually reduce back to the `step_val + 3 == step_grad_val`
      # case.
  
      # Logic as above: in these cases we don't need to store a checkpoint at all.
      no_checkpoint = index == checkpoints
      no_checkpoint = no_checkpoint | (step_val + 1 == step_grad_val)
      no_checkpoint = no_checkpoint | (step_val + 2 == step_grad_val)
  
      step_next_checkpoint = jnp.where(
          no_checkpoint, _unreachable_checkpoint_step(out), out
      )
      # Invariant: `step_val < step_next_checkpoint`. (Due to the invariant
      # `step_val + 1 <= step_grad_val`.)
      step_next_checkpoint = nonbatchable(step_next_checkpoint)
      return step_next_checkpoint
  
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 18 for calling a global function with example usage: _calc_next_checkpoint(step_val, step_grad_val, index, checkpoints) and returns: step_next_checkpoint