code: def test_nested_loops(read, getkey):
      @ft.partial(jax.jit, static_argnums=5)
      @ft.partial(jax.vmap, in_axes=(0, 0, 0, 0, 0, None))
      def run(step, vals, ts, final_step, cotangents, true):
          val1, val2, val3, val4 = vals
          value, vjp_fn = jax.vjp(
              lambda _val1: outer_loop(
                  step, (_val1, val2, val3, val4), ts, true, final_step
              ),
              val1,
          )
          cotangents = vjp_fn(cotangents)
          return value, cotangents
  
      def outer_loop(step, vals, ts, true, final_step):
          def cond(carry):
              step, _ = carry
              return step < final_step
  
          def body(carry):
              step, (val1, val2, val3, val4) = carry
              mul = 1 + 0.05 * jnp.sin(105 * val1 + 1)
              val1 = val1 * mul
              return inner_loop(
                  step, (val1, val2, val3, val4), ts, true
              )
  
          def buffers(carry):
              _, (_, val2, val3, _) = carry
              return val2, val3
  
          if true:
              while_loop = ft.partial(_while_as_scan, max_steps=50)
          else:
              while_loop = ft.partial(
                  eqxi.while_loop,
                  max_steps=50,
                  buffers=buffers,
                  kind="checkpointed",
              )
          _, out = while_loop(cond, body, (step, vals))
          return out
  
      def inner_loop(step, vals, ts, true):
          ts_done = jnp.floor(ts[step] + 1)
  
          def cond(carry):
              step, _ = carry
              return ts[step] < ts_done
  
          def body(carry):
              step, (val1, val2, val3, val4) = carry
              mul = 1 + 0.05 * jnp.sin(100 * val1 + 3)
              val1 = val1 * mul
              if read:
                  i2 = jr.randint(
                      jr.PRNGKey(step), shape=(), minval=0, maxval=step
                  )
                  i3 = jr.randint(
                      jr.PRNGKey(step + 100),
                      shape=(),
                      minval=0,
                      maxval=step,
                  )
                  i4 = jr.randint(
                      jr.PRNGKey(step + 200),
                      shape=(),
                      minval=0,
                      maxval=step,
                  )
                  setval2 = val1 + val2[i2]
                  setval3 = 0.8 * val1 + val2[i3] + val3[i3]
                  setval4 = val1 + jnp.tanh(val2[i2] * val4[i4])
                  setval2 = jnp.where(step == 0, 2.0, setval2)
                  setval3 = jnp.where(step == 0, 3.0, setval3)
                  setval4 = jnp.where(step == 0, 4.0, setval4)
              else:
                  setval2 = val1
                  setval3 = val1
                  setval4 = val1
              val2 = val2.at[step].set(setval2)
              val3 = val3.at[step].set(setval3)
              val4 = val4.at[step].set(setval4)
              return step + 1, (val1, val2, val3, val4)
  
          def buffers(carry):
              _, (_, _, val3, val4) = carry
              return val3, val4
  
          if true:
              while_loop = ft.partial(_while_as_scan, max_steps=10)
          else:
              while_loop = ft.partial(
                  eqxi.while_loop,
                  max_steps=10,
                  buffers=buffers,
                  kind="checkpointed",
              )
          return while_loop(cond, body, (step, vals))
  
      step = jnp.array([0, 5])
      val1 = jr.uniform(getkey(), shape=(2,), minval=0.1, maxval=0.7)
      val2 = val3 = val4 = jnp.zeros((2, 47))
      ts = jnp.stack([jnp.linspace(0, 19, 47), jnp.linspace(0, 13, 47)])
      final_step = jnp.array([46, 43])
      cotangents = (
          jr.normal(getkey(), (2,)),
          jr.normal(getkey(), (2, 47)),
          jr.normal(getkey(), (2, 47)),
          jr.normal(getkey(), (2, 47)),
      )
  
      value, grads = run(
          step,
          (val1, val2, val3, val4),
          ts,
          final_step,
          cotangents,
          False,
      )
      true_value, true_grads = run(
          step,
          (val1, val2, val3, val4),
          ts,
          final_step,
          cotangents,
          True,
      )
  
      assert tree_allclose(value, true_value)
      assert tree_allclose(grads, true_grads, rtol=1e-4, atol=1e-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 14 for calling a global function with example usage: test_nested_loops(read, getkey) and returns: value,  cotangents, out, while_loop, unknown, inner_loop, val2,  val3, unknown, unknown,  val1,  val2,  val3,  val4, val3,  val4