code: def test_speed_while(inplace_op, while_loop):
      @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 inplace_op == "scatter":
                  xs = xs.at[step].set(1)
              elif inplace_op == "dynamic_update_slice":
                  xs = lax.dynamic_update_index_in_dim(xs, 1.0, step, 0)
              else:
                  assert False
              return step + 1, xs
  
          return while_loop(cond, body, (init_step, init_xs))
  
      size = 100_000
      args = jnp.array([0]), jnp.zeros((1, size))
      f(*args)  # compile
  
      speed = timeit.timeit(lambda: f(*args), number=1)
      # Takes O(1e-3) with optimisation.
      # Takes O(10) without optimisation.
      # So we have two orders of magnitude safety margin each way, so the test shouldn't
      # be flaky.
      assert speed < 0.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 11 for calling a global function with example usage: test_speed_while(inplace_op, while_loop) and returns: while_loop, unknown, unknown,  xs