code: def test_forward(buffer, kind, getkey):
      cond_fun, make_body_fun, init_val1, init_val2, mlp = _get_problem(
          getkey(), num_steps=5
      )
      body_fun = make_body_fun(mlp)
      true_final_carry = lax.while_loop(
          cond_fun, body_fun, (0, init_val1, init_val2)
      )
  
      @jax.jit
      def run(init_val1, init_val2):
          if buffer:
              buffer_fn = lambda i: i[2]
          else:
              buffer_fn = None
          return eqxi.while_loop(
              cond_fun,
              body_fun,
              (0, init_val1, init_val2),
              max_steps=10 if kind == "bounded" else None,
              buffers=buffer_fn,
              kind=kind,
              checkpoints=9,
          )
  
      final_carry, _ = jax.linearize(run, init_val1, init_val2)
      assert tree_allclose(
          final_carry, true_final_carry, atol=1e-4, rtol=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 4 for calling a global function with example usage: test_forward(buffer, kind, getkey) and returns: eqxi