code: def test_speed_grad_checkpointed_while(getkey):
      mlp = eqx.nn.MLP(2, 1, 2, 2, key=getkey())
      checkpoints = 10_000
  
      @jax.jit
      @jax.vmap
      @jax.grad
      def f(init_val, init_step):
          def cond(carry):
              step, _ = carry
              return step < 2 * checkpoints
  
          def body(carry):
              step, val = carry
              (theta,) = mlp(val)
              real, imag = val
              z = (
                  real.astype(jnp.complex64)
                  + imag.astype(jnp.complex64) * 1j
              )
              z = z * jnp.exp(1j * theta.astype(jnp.complex64))
              real = jnp.real(z)
              imag = jnp.imag(z)
              return step + 1, jnp.stack([real, imag])
  
          _, final_xs = eqxi.while_loop(
              cond,
              body,
              (init_step, init_val),
              kind="checkpointed",
              checkpoints=checkpoints,
          )
          return jnp.sum(final_xs)
  
      init_step = jnp.array([0, 10])
      init_val = jr.normal(getkey(), (2, 2))
  
      f(init_val, init_step)  # compile
      speed = timeit.timeit(lambda: f(init_val, init_step), number=1)
      # Should take ~0.01 seconds
      assert speed < 0.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 13 for calling a global function with example usage: test_speed_grad_checkpointed_while(getkey) and returns: jnp, unknown, unknown,  jnp