code: def test_buffer_index():
      def cond_fun(carry):
          return True
  
      def body_fun(carry):
          return carry.at[..., 1:].set(0)
  
      def buffers(carry):
          return carry
  
      init = jnp.zeros((3, 5))
      eqxi.while_loop(
          cond_fun,
          body_fun,
          init,
          kind="checkpointed",
          buffers=buffers,
          max_steps=2,
      )
  
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 18 for calling a global function with example usage: test_buffer_index() and returns: True, carry, carry