code: def test_backward_checkpointed(
      checkpoints,
      num_steps,
      backward_order,
      with_max_steps,
      buffer,
      getkey,
      capfd,
  ):
      if with_max_steps:
          max_steps = num_steps
          get_num_steps = None
      else:
          max_steps = None
          get_num_steps = num_steps
      cond_fun, make_body_fun, init_val1, init_val2, mlp = _get_problem(
          getkey(), num_steps=get_num_steps
      )
  
      @jax.jit
      @jax.value_and_grad
      def true_run(arg):
          init_val1, init_val2, mlp = arg
          body_fun = make_body_fun(mlp)
          _, true_final_val1, true_final_val2 = _while_as_scan(
              cond_fun,
              body_fun,
              (0, init_val1, init_val2),
              max_steps=num_steps,
          )
          return jnp.sum(true_final_val1) + jnp.sum(true_final_val2)
  
      @jax.jit
      @jax.value_and_grad
      def run(arg):
          init_val1, init_val2, mlp = arg
          if buffer:
              buffer_fn = lambda i: i[2]
          else:
              buffer_fn = None
          body_fun = make_body_fun(mlp)
          _, final_val1, final_val2 = eqxi.while_loop(
              cond_fun,
              body_fun,
              (0, init_val1, init_val2),
              max_steps=max_steps,
              buffers=buffer_fn,
              kind="checkpointed",
              checkpoints=checkpoints,
          )
          return jnp.sum(final_val1) + jnp.sum(final_val2)
  
      true_value, true_grad = true_run([init_val1, init_val2, mlp])
      capfd.readouterr()
      value, grad = run([init_val1, init_val2, mlp])
      text, _ = capfd.readouterr()
      true_text = "".join(
          f"{i}\n" for i in range(num_steps)
      ) + backward_order.replace(",", "\n")
      if buffer:
          grad[1] = grad[1].at[:num_steps].set(0)
      assert tree_allclose(value, true_value, rtol=1e-4, atol=1e-4)
      assert tree_allclose(grad, true_grad, rtol=1e-4, atol=1e-4)
      assert true_text in text.strip()
  
imports: import functools
  import timeit
  from typing import Optional
  import equinox
  import equinox.internal
  import jax
  import jax.lax
  import jax.lib
  import jax.numpy
  import jax.random
  import jax.tree_util
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 5 for calling a global function with example usage: test_backward_checkpointed(checkpoints, num_steps, backward_order, with_max_steps, buffer, getkey, capfd) and returns: unknown, unknown