code: def test_vmap_primal_batched_cond(buffer, kind, getkey):
      cond_fun, make_body_fun, init_val1, init_val2, mlp = _get_problem(
          getkey(), num_steps=14
      )
  
      value_and_grad = _maybe_value_and_grad(kind)
  
      @jax.jit
      @ft.partial(jax.vmap, in_axes=([0, 0, None], 0))
      @value_and_grad
      def true_run(arg, init_step):
          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,
              (init_step, init_val1, init_val2),
              max_steps=14,
          )
          return jnp.sum(true_final_val1) + jnp.sum(true_final_val2)
  
      @jax.jit
      @ft.partial(jax.vmap, in_axes=([0, 0, None], 0))
      @value_and_grad
      def run(arg, init_step):
          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,
              (init_step, init_val1, init_val2),
              max_steps=16 if kind == "bounded" else None,
              buffers=buffer_fn,
              kind=kind,
              checkpoints=4,
          )
          return jnp.sum(final_val1) + jnp.sum(final_val2)
  
      init_step = jnp.array([0, 1, 2, 3, 5, 10])
      init_val1, init_val2 = jtu.tree_map(
          lambda x: jr.normal(getkey(), (6,) + x.shape, x.dtype),
          (init_val1, init_val2),
      )
      true_value, true_grad = true_run(
          [init_val1, init_val2, mlp], init_step
      )
      value, grad = run([init_val1, init_val2, mlp], init_step)
      if buffer and kind != "lax":
          for i, j in enumerate(init_step):
              grad[1] = grad[1].at[i, j.item() : 14].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)
  
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 9 for calling a global function with example usage: test_vmap_primal_batched_cond(buffer, kind, getkey) and returns: unknown, unknown