code: def while_stmt(test, body, get_state, set_state, symbol_names, opts):
      """Functional form of a while statement.
  
      The loop operates on a so-called state, which includes all symbols that are
      variant across loop iterations. In what follows we refer to state as either
      a tuple of entities that represent an actual state, or a list of arguments
      of the corresponding types.
  
      The inputs and outputs of the callables representing the loop blocks are not
      explicit - instead, these functions must use nonlocal/global for side effects.
      The inputs and outputs are instead controlled by the set_state/get_state
      functions.
  
      Args:
        test: Callable with boolean return type. The loop condition.
        body: Callable representing the actual loop body.
        get_state: Additional callable which can capture additional state (such as
          the values of composite symbols). This is only useful when staging the
          loop.
        set_state: Additional callable which save values captured by get_state back
          into the Python environment. This is only useful when staging the loop.
        symbol_names: Tuple containing the names of all loop variables.
        opts: Optional dict of extra loop parameters.
  
      Returns:
        Tuple containing the final state.
      """
  
      # Evaluate the initial test once in order to do the dispatch. The evaluation
      # is isolated to minimize unwanted side effects.
      # TODO(mdan): Do a full iteration - some state types might lower to Tensor.
      with func_graph.FuncGraph("tmp").as_default():
          init_test = test()
  
      # TensorFlow: Multiple evaluations are acceptable in this case, so we're fine
      # with the re-evaluation of `test` that `_tf_while_stmt` will make.
      if tensors.is_dense_tensor(init_test):
          _tf_while_stmt(
              test, body, get_state, set_state, symbol_names, opts
          )
          return
  
      # Normal Python: We already consumed one evaluation of `test`; consistently,
      # unroll one iteration before dispatching to a normal loop.
      # TODO(mdan): Push the "init_test" value via opts into _py_while_stmt?
      if not init_test:
          return
      body()
  
      _py_while_stmt(test, body, get_state, set_state, opts)
  
imports: import functools
  import sys
  import traceback
  import numpy
  from tensorflow.python.autograph.operators import py_builtins
  from tensorflow.python.autograph.operators import variables
  from tensorflow.python.autograph.utils import ag_logging
  from tensorflow.python.autograph.utils import misc
  from tensorflow.python.autograph.utils import tensors
  from tensorflow.python.autograph.utils import type_registry
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import func_graph
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_conversion
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import cond
  from tensorflow.python.ops import control_flow_assert
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import tensor_array_ops
  from tensorflow.python.ops import while_loop
  from tensorflow.python.types import distribute
  from tensorflow.python.util import nest
  from tensorflow.python.util import variable_utils
label_desc: Use image node_id 17 for calling a global function with example usage: while_stmt(test, body, get_state, set_state, symbol_names, opts) without return types