code: def _while_as_scan(cond, body, init_val, max_steps):
      def f(val, _):
          val2 = lax.cond(cond(val), body, lambda x: x, val)
          return val2, None
  
      final_val, _ = lax.scan(f, init_val, xs=None, length=max_steps)
      return final_val
  
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 2 for calling a global function with example usage: _while_as_scan(cond, body, init_val, max_steps) and returns: final_val, val2,  None