code: def test_vmap_cotangent(buffer, kind, getkey):
      cond_fun, make_body_fun, init_val1, init_val2, mlp = _get_problem(
          getkey(), num_steps=14
      )
  
      @jax.jit
      @jax.jacrev
      def true_run(arg, init_val2):
          init_val1, 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=14,
          )
          return true_final_val1, true_final_val2
  
      @jax.jit
      @jax.jacrev
      def run(arg, init_val2):
          init_val1, 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=16 if kind == "bounded" else None,
              buffers=buffer_fn,
              kind=kind,
              checkpoints=4,
          )
          return final_val1, final_val2
  
      true_jac = true_run((init_val1, mlp), init_val2)
      jac = run((init_val1, mlp), init_val2)
      assert tree_allclose(jac, true_jac, rtol=1e-4, atol=1e-4)
  
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 10 for calling a global function with example usage: test_vmap_cotangent(buffer, kind, getkey) and returns: true_final_val1,  true_final_val2, final_val1,  final_val2