code: def test_speed_buffer_while(read):
      @jax.jit
      @jax.vmap
      def f(init_step, init_xs):
          def cond(carry):
              step, xs = carry
              return step < xs.size
  
          def body(carry):
              step, xs = carry
              if read:
                  i = jr.randint(
                      jr.PRNGKey(step), shape=(), minval=0, maxval=step
                  )
                  xs = xs.at[step].set(2 * jnp.tanh(xs[i]))
              else:
                  xs = xs.at[step].set(1)
              return step + 1, xs
  
          def loop(init_xs):
              return eqxi.while_loop(
                  cond,
                  body,
                  (init_step, init_xs),
                  buffers=lambda i: i[1],
                  kind="checkpointed",
                  checkpoints=50_000,
              )
  
          # Linearize so that we save residuals
          return jax.linearize(loop, init_xs)
  
      size = 100_000
      # nontrivial batch size is important to ensure that the `.at[].set()` is really a
      # scatter, and that XLA doesn't optimise it into a dynamic_update_slice. (Which
      # can be switched with `select` in the compiler.)
      args = jnp.array([0, 1]), jnp.zeros((2, size))
      f(*args)  # compile
  
      speed = timeit.timeit(lambda: f(*args), number=1)
      assert speed < 1
  
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 12 for calling a global function with example usage: test_speed_buffer_while(read) and returns: jax, unknown, unknown,  xs, eqxi