code: def test_notangent_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)
      )
      if buffer:
          buffer_fn = lambda i: i[2]
      else:
          buffer_fn = None
      final_carry = eqxi.while_loop(
          cond_fun,
          body_fun,
          (0, init_val1, init_val2),
          max_steps=10 if kind == "bounded" else None,
          kind=kind,
          buffers=buffer_fn,
          checkpoints=1,
      )
      assert tree_allclose(
          final_carry, true_final_carry, rtol=1e-5, atol=1e-5
      )
  
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 3 for calling a global function with example usage: test_notangent_forward(buffer, kind, getkey) without return types