code: def _should_save_residual(
      step,
      save_state,
      residual_steps,
      u2_minus_1,
      checkpoints,
      max_steps,
  ):
      """This is the controller for whether we should save the current value at each step,
      and if so which memory location to save it in.
      """
      # TODO: also implement Algorithm 2 of Stumm and Walther, which gives improved
      # results for u2 < step < u3.
  
      if checkpoints == max_steps:
          # Important special case! We don't need to do any computations in this case.
          # This is a measurable performance optimisation.
          save_residual = True
          index = step
          save_state2 = save_state
      else:
          if max_steps is not None and max_steps <= u2_minus_1:
              # `step < max_steps` implies `step + 1 <= max_steps`. If it is also the case
              # that `max_steps <= u2 - 1` then overall we know that `step <= u2 - 2`,
              # which is the condition for using Stumm--Walther.
              #
              # Skipping this cond offers a measurable performance optimisation.
              pred = False
          else:
              step, u2_minus_1 = nonbatchable((step, u2_minus_1))
              pred = step >= u2_minus_1
          save_residual, index, save_state2 = lax.cond(
              pred,
              _wang_moin_wrapper,
              _stumm_walther_i_wrapper,
              step,
              save_state,
              residual_steps,
          )
      return save_residual, index, save_state2
  
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 12 for calling a global function with example usage: _should_save_residual(step, save_state, residual_steps, u2_minus_1, checkpoints, max_steps) and returns: save_residual,  index,  save_state2