code: def _checkpointed_while_loop_bwd(
      remainders,
      grad_final_val,
      perturbed,
      vjp_arg,
      cond_fun,
      checkpoints,
      buffers,
      max_steps,
  ):
      """Time for the complicated bit: iterate backward through a checkpointed while loop,
      loading values from checkpoints and using treeverse to toggle between forward and
      backward steps.
      """
      perturb_init_val, perturb_body_fun = perturbed
      _, body_fun = vjp_arg
      grad_final_body_fun = jtu.tree_map(
          lambda x, p: jnp.zeros_like(x) if p else None,
          body_fun,
          perturb_body_fun,
      )
      del cond_fun, perturbed
      (
          num_steps,
          init_residual_steps,
          init_residuals,
          filled_buffers,
      ) = remainders
      num_steps, init_residual_steps = nonbatchable(
          (num_steps, init_residual_steps)
      )
  
      def _cond_fun(carry):
          _, step_grad_val, *_ = carry
          # step_grad_val is the location of our cotangent. We want to keep going until
          # this has got all the way to the start.
          step_grad_val = nonbatchable(step_grad_val)
          return step_grad_val > 0
  
      def _body_fun(carry):
          (
              step_val,
              step_grad_val,
              step_next_checkpoint,
              index,
              val,
              grad_val,
              grad_body_fun,
              residual_steps,
              residuals,
          ) = carry
          (
              step_val,
              step_grad_val,
              step_next_checkpoint,
              index,
              residual_steps,
          ) = nonbatchable(
              (
                  step_val,
                  step_grad_val,
                  step_next_checkpoint,
                  index,
                  residual_steps,
              )
          )
  
          msg = (
              "Internal run-time error when backpropagating through "
              "`equinox.internal.checkpointed_while_loop`. "
              "Please raise an issue at https://github.com/patrick-kidger/equinox"
          )
          if getattr(typing, "TESTING", False):
              step_val = error_if(
                  step_val, step_val >= step_grad_val, msg
              )
  
          #
          # First either propagate our primal state forward, or make a U-turn if the
          # primal state has caught up to the cotangent state.
          #
  
          # Use `filter_vjp` to neatly handle floating-point arrays.
          #
          # We pass in `body_fun` as an argument as it contains its closed-over values
          # in its PyTree structure, and we do want to compute cotangents wrt these.
          val_buffers = tree_at(
              buffers(_is_none), val, filled_buffers, is_leaf=_is_none
          )
          grad_val_has_ct = jtu.tree_map(
              is_array, grad_val, is_leaf=_is_none
          )
          body_fun_has_ct = jtu.tree_map(
              is_array, grad_body_fun, is_leaf=_is_none
          )
          diff_val, nondiff_val = partition(
              val_buffers, grad_val_has_ct
          )
          diff_body, nondiff_body = partition(body_fun, body_fun_has_ct)
  
          def _to_vjp(_diff_body, _diff_val):
              _val_buffers = combine(_diff_val, nondiff_val)
              _body_fun = combine(_diff_body, nondiff_body)
              _out = _body_fun(_val_buffers)
              return partition(_out, grad_val_has_ct)
  
          dynamic_val_candidate, vjp_fn, static_val_candidate = jax.vjp(
              _to_vjp, diff_body, diff_val, has_aux=True
          )
          val_candidate = combine(
              dynamic_val_candidate, static_val_candidate
          )
          val_candidate = tree_at(
              buffers(None), val_candidate, replace_fn=lambda _: None
          )
  
          if checkpoints == max_steps:
              # Skip calculating this at runtime; we know it must be true.
              # This means we can inline the U-turn branch of the `cond`.
              perform_u_turn = True
          else:
              perform_u_turn = step_val + 1 == step_grad_val
              perform_u_turn = nonbatchable(perform_u_turn)
          (
              step_val2,
              step_grad_val2,
              step_next_checkpoint2,
              index2,
              val2,
              grad_val2,
              grad_body_fun2,
          ) = lax.cond(
              perform_u_turn,
              _make_u_turn(
                  vjp_fn,
                  residual_steps,
                  residuals,
                  checkpoints,
                  max_steps,
              ),
              _fwd,
              step_val,
              step_grad_val,
              step_next_checkpoint,
              index,
              val_candidate,
              grad_val,
              grad_body_fun,
          )
  
          #
          # Second, decide whether to store our current primal state in a checkpoint.
          # Note that this can only actually trigger on `_fwd` and not on `_u_turn`, as
          # if `_u_turn` happens then `step_val2 < step_next_checkpoint2`, but
          # `_maybe_save_to_checkpoint` has a `step_val2 == step_next_checkpoint2` check.
          # (We could maybe move this inside `_fwd`? It was originally outside for
          # efficiency reasons that I don't think apply any more.)
          #
  
          if checkpoints == max_steps:
              # In this case we never need to save any extra checkpoints on the backward
              # pass, so skip this.
              residual_steps2 = residual_steps
              residuals2 = residuals
          else:
              (
                  index2,
                  step_next_checkpoint2,
                  residual_steps2,
                  residuals2,
              ) = _maybe_save_to_checkpoint(
                  step_val2,
                  step_grad_val2,
                  step_next_checkpoint2,
                  index2,
                  val2,
                  residual_steps,
                  residuals,
                  checkpoints,
              )
  
          return (
              step_val2,
              step_grad_val2,
              step_next_checkpoint2,
              index2,
              val2,
              grad_val2,
              grad_body_fun2,
              residual_steps2,
              residuals2,
          )
  
      # We can index into our residuals using 0, 1, ..., checkpoints - 1.
      # `index` is used to refer to the next empty spot, so it takes values in
      # 0, 1, ..., checkpoints - 1, checkpoints, where `index == checkpoints` indicates
      # that there are no empty spots and the whole buffer is full. (And this is used in
      # `_calc_step_next_checkpoint`.)
      init_index = jnp.minimum(num_steps, checkpoints)
      init_step_grad_val = num_steps
      init_step_val, init_val, init_index = _load_from_checkpoint(
          init_step_grad_val,
          init_index,
          init_residual_steps,
          init_residuals,
          checkpoints,
          max_steps,
      )
      init_step_next_checkpoint = _calc_next_checkpoint(
          init_step_val, init_step_grad_val, init_index, checkpoints
      )
  
      #
      # Okay, this next bit is a bit complicated.
      # We need to figure out exactly which parts of the carry that we're going to compute
      # cotangents for.
      # There are three interacting pieces:
      # (1) Only some inputs even need their cotangents computed (those that are marked
      #     as perturbed). We should avoid computing intermediate cotangents that aren't
      #     needed.
      # (2) Some cotangents may be symbolic zeros. We should avoid materialising any that
      #     don't need to be materialised.
      # (3) Some parts of the body function may not be differentiable. (E.g. they call to
      #     `eqxi.nondifferentiable` or `eqxi.nondifferentiable_backward`.) In particular,
      #     some inexact arrays may be marked with these. (Anything else is automatically
      #     nondifferentiable.) As these raise an error eagerly, we can't just take an
      #     approach of "differentiate everything, then DCE".
      #
      # Our strategy is as follows. First, consider those inputs that are perturbed, and
      # iterate to find all the carries that become perturbed as a result. We obtain a
      # fixed point.
      # * This mimicks what JAX does in JVP rules for scans/whiles. And in fact we also
      #   track perturbation information forward via JVPs (which input tangents affect
      #   which output tangents), rather than backward via VJPs (which output cotangents
      #   affect which input cotangents), as this is what JAX does here and we'd like to
      #   be consistent. (Also, it seems like there might be a catch-22 trying to do this
      #   with VJPs -- how do you decide which inputs to perturb in your VJP?) This is
      #   basically handling criteria (1).
      # * This also means that we're perturbing the minmal amount possible. We don't even
      #   run the JVP rule for something that isn't differentiated. This is part of
      #   addressing criteria (3).
      #
      # Second, we ignore all cotangents (of our forward output) that aren't on a
      # perturbed carry. Clearly they're not going to affect anything. So we switch them
      # to symbolic zeros.
      # * This is part of addressing criteria (3) -- we avoid propagating cotangents
      #   that we don't need to. Maybe there's an `eqxi.nondifferentiable_backward` in
      #   there.
      #
      # Third, we find another fixed point, this time for the symbolic zeros. Track which
      # cotangents affect which cotangents, and materialise the minimal amount possible.
      # * This addresses criteria (2).
      # * Note that symbolic zeros are doing double-duty here. We avoid materialising them
      #   for simple efficiency reasons, sure -- but we also want to keep as many
      #   unmaterialised as possible, in case there's an `eqxi.nondifferentiable_backward`
      #   in there, which only accepts symbolic zero cotangents.
      #
      # Phew! At that point, job done: we know which cotangents we're propagating and so
      # we can run our loop.
      #
      # The first fixed point basically mimicks what JAX does in the JVP rule for a while
      # or scan, figuring out which values are perturbed. The second fixed point basically
      # mimicks what JAX does in the transpose rule of a scan, figuring out which
      # cotangents are symbolic zeros.
      #
  
      # Do this inside the dynamic context of a `jax.eval_shape` so that it's as cheap as
      # possible, without extra tracing.
      def _resolve_perturb_val():
          init_val_buffers = tree_at(
              buffers(_is_none),
              init_val,
              filled_buffers,
              is_leaf=_is_none,
          )
          perturb_val = perturb_init_val
          assert jtu.tree_structure(
              init_val_buffers
          ) == jtu.tree_structure(perturb_val)
  
          while True:
              # `body_fun` is included so that we also track perturbatations on that.
              perturb_pair = (perturb_body_fun, perturb_val)
              dynamic, static = partition(
                  (body_fun, init_val_buffers), perturb_pair
              )
              new_perturb_val = sentinel = object()
  
              @jax.custom_jvp
              def _record_symbolic_zeros(_dynamic_out):
                  return _dynamic_out
  
              def _record_symbolic_zeros_jvp(primals, tangents):
                  (primals,) = primals
                  (tangents,) = tangents
                  nonlocal new_perturb_val
                  new_perturb_val = jtu.tree_map(
                      _not_symbolic_zero, tangents
                  )
                  return primals, tangents
  
              _record_symbolic_zeros.defjvp(
                  _record_symbolic_zeros_jvp, symbolic_zeros=True
              )
  
              def _to_linearize(_dynamic):
                  _body_fun, _val = combine(_dynamic, static)
                  _out = _body_fun(_val)
                  _dynamic_out, _static_out = partition(
                      _out, is_inexact_array
                  )
                  _dynamic_out = _record_symbolic_zeros(_dynamic_out)
                  _out = combine(_dynamic_out, _static_out)
                  return _out
  
              # Not `jax.jvp`, so as not to error if `body_fun` has any `custom_vjp`s.
              jax.linearize(_to_linearize, dynamic)
              assert new_perturb_val is not sentinel
              assert jtu.tree_structure(
                  perturb_val, is_leaf=_is_none
              ) == jtu.tree_structure(new_perturb_val, is_leaf=_is_none)
              new_perturb_val = jtu.tree_map(
                  lambda x, y: False
                  if (x is False and y is None)
                  else y,
                  perturb_val,
                  new_perturb_val,
              )
              assert jtu.tree_structure(
                  perturb_val
              ) == jtu.tree_structure(new_perturb_val)
              if tree_equal(perturb_val, new_perturb_val):
                  jtu.tree_map(_assert_bool, perturb_val)
                  if getattr(typing, "TESTING", False):
                      print("perturb_val", perturb_val)
                  return Static(perturb_val)
              else:
                  perturb_val = jtu.tree_map(
                      operator.or_, perturb_val, new_perturb_val
                  )
  
      # Find a fixed point of which values we need to perturb. (Not the same as whether
      # or not they have symbolic zero cotangents! All unperturbed values must have
      # symbolic zero cotangents, but some perturbed values may have these as well.)
      perturb_val = jax.eval_shape(_resolve_perturb_val).value
  
      # Do this inside the dynamic context of a `jax.eval_shape` so that it's as cheap as
      # possible, without extra tracing.
      def _resolve_symbolic_zeros(grad_val):
          symbolic_zero_gradient = jtu.tree_map(
              _is_none, grad_val, is_leaf=_is_none
          )
          init_val_buffers = tree_at(
              buffers(_is_none),
              init_val,
              filled_buffers,
              is_leaf=_is_none,
          )
          dynamic_init_val, static_init_val = partition(
              init_val_buffers, perturb_val
          )
  
          while True:
              new_symbolic_zero_gradient = sentinel = object()
  
              # Some hackery to extract which inputs still have symbolic zero cotangents.
              @jax.custom_vjp
              def _record_symbolic_zero(_dynamic_val):
                  return _dynamic_val
  
              def _record_symbolic_zero_fwd(_dynamic_val):
                  return jtu.tree_map(_get_value, _dynamic_val), None
  
              def _record_symbolic_zero_bwd(_, grad_dynamic_val):
                  nonlocal new_symbolic_zero_gradient
                  new_symbolic_zero_gradient = jtu.tree_map(
                      _is_symbolic_zero_or_none,
                      grad_dynamic_val,
                      is_leaf=_is_none,
                  )
                  return (grad_dynamic_val,)
  
              _record_symbolic_zero.defvjp(
                  _record_symbolic_zero_fwd,
                  _record_symbolic_zero_bwd,
                  symbolic_zeros=True,
              )
  
              def _to_vjp(_dynamic_val):
                  _dynamic_val = _record_symbolic_zero(_dynamic_val)
                  val = combine(_dynamic_val, static_init_val)
                  return filter(
                      body_fun(val),
                      symbolic_zero_gradient,
                      inverse=True,
                  )
  
              _, vjp_fn = jax.vjp(_to_vjp, dynamic_init_val)
              vjp_fn(grad_val)
              assert new_symbolic_zero_gradient is not sentinel
              if tree_equal(
                  symbolic_zero_gradient, new_symbolic_zero_gradient
              ):
                  jtu.tree_map(_assert_bool, symbolic_zero_gradient)
                  if getattr(typing, "TESTING", False):
                      print(
                          "symbolic_zero_gradient",
                          symbolic_zero_gradient,
                      )
                  return Static(symbolic_zero_gradient)
              else:
                  symbolic_zero_gradient = jtu.tree_map(
                      operator.and_,
                      symbolic_zero_gradient,
                      new_symbolic_zero_gradient,
                  )
                  grad_val = jtu.tree_map(
                      _materialise_symbolic_zero,
                      symbolic_zero_gradient,
                      grad_val,
                      init_val_buffers,
                  )
  
      grad_final_val = filter(grad_final_val, perturb_val)
      # If all provided cotangents end up being irrelevant -- i.e. the perturbed inputs
      # depend only on those outputs which have symbolic zero cotangents -- then we can
      # skip the whole computation.
      if len(jtu.tree_leaves(grad_final_val)) == 0:
          return jtu.tree_map(
              lambda _: None, (grad_final_val, body_fun)
          )
      symbolic_zero_gradient = jax.eval_shape(
          _resolve_symbolic_zeros, grad_final_val
      ).value
      init_val_buffers = tree_at(
          buffers(_is_none), init_val, filled_buffers, is_leaf=_is_none
      )
      grad_final_val = jtu.tree_map(
          _materialise_symbolic_zero,
          symbolic_zero_gradient,
          grad_final_val,
          init_val_buffers,
      )
  
      init_carry = (
          init_step_val,
          init_step_grad_val,
          init_step_next_checkpoint,
          init_index,
          init_val,
          grad_final_val,
          grad_final_body_fun,
          init_residual_steps,
          init_residuals,
      )
      # Note that the saved checkpoints hold both (a) values computed on the forward
      # pass, and (b) checkpoints recomputed on the backward pass. (We don't need to
      # distinguish them.)
      #
      # Controller State
      # ----------------
      # `step_val`: see `val`.
      # `step_grad_val`: see `grad_val`. Note that there is an invariant
      #   `step_val + 1 <= step_grad_val`.
      # `step_next_checkpoint`: the step at which we next need to save a checkpoint, for
      #   the recomputed forward computation.
      # `index` is the index of the memory buffer to save the next checkpoint, for the
      #   recomputed forward computation.
      #
      # Numerical computations
      # ----------------------
      # `val` is the evolving state of the forward loop (reloaded from checkpoints). The
      #   forward loop step that this comes from is `step_val`. As such `val` and
      #   `step_val` will jump back-and-forth as values are loaded from checkpoints and
      #   then recomputed forward using `_fwd`.
      # `grad_val` is the cotangent that we're propagating backward. The step of the
      #   forward loop (that it holds the cotangent for) is given in `step_grad_val`. As
      #   such `step_grad_val` decrements weakly monotonically, and `grad_val` is updated
      #   on every U-turn.
      # `grad_body_fun` is the cotangent being accumulated for `body_fun`. It updates on
      #   every U-turn.
      #
      # Checkpoints
      # -------
      # `residual_steps` is the memory holding the `step` for each checkpoint
      # `residuals` is the memory holding the `val` for each checkpoint
  
      # Not reverse mode autodifferentiable, but it is forward-mode autodifferentiable!
      # That means we can do Hessians, at least.
      while_loop = jax.named_call(
          lax.while_loop, name="checkpointed-bwd"
      )
      final_carry = while_loop(_cond_fun, _body_fun, init_carry)
      *_, grad_init_val, grad_body_fun, _, _ = final_carry
      out = grad_init_val, grad_body_fun
      return out
  
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 27 for calling a global function with example usage: _checkpointed_while_loop_bwd(remainders, grad_final_val, perturbed, vjp_arg, cond_fun, checkpoints, buffers, max_steps) and returns: out, unknown, step_val2,  step_grad_val2,  step_next_checkpoint2,  index2,  val2,  grad_val2,  grad_body_fun2,  residual_steps2,  residuals2, jtu, partition, _dynamic_out, primals,  tangents, _out, Static, _dynamic_val, jtu,  None, grad_dynamic_val, filter, Static