code: def while_loop(test_fn, body_fn, vars):
      def body_fn_wrapper(loop_vars):
          return body_fn(*loop_vars)
  
      def test_fn_wrapper(loop_vars):
          return test_fn(*loop_vars)
  
      if isinstance(vars, dict):
          vars = list(vars.values())
      with jax.disable_jit():
          final_loop_vars = jax.lax.while_loop(
              test_fn_wrapper, body_fn_wrapper, vars
          )
      return final_loop_vars
  
label_desc: Use image node_id 2 for calling a global function with example usage: while_loop(test_fn, body_fn, vars) and returns: final_loop_vars, body_fn, test_fn