code: def while_loop(
      cond, func, loop_vars, max_iterations=None, name="while_loop"
  ):
      """Run a while loop with user-defined computation and loop condition.
  
      This operator simulates a while loop which iterately does customized computation
      as long as the condition is satisfied.
  
      `loop_vars` is a list of NDArrays on which the computation uses.
  
      `cond` is a user-defined function, used as the loop condition.
      It consumes `loop_vars`, and produces a scalar MXNet NDArray,
      indicating the termination of the loop.
      The loop ends when `cond` returns false (zero).
      The `cond` is variadic, and its signature should be
      `cond(*loop_vars) => NDArray`.
  
      `func` is a user-defined function, used as the loop body.
      It also consumes `loop_vars`, and produces `step_output` and `new_loop_vars` at each step.
      In each step, `step_output` should contain the same number elements.
      Through all steps, the i-th element of `step_output` should have the same shape and dtype.
      Also, `new_loop_vars` should contain the same number of elements as `loop_vars`,
      and the corresponding element should have the same shape and dtype.
      The `func` is variadic, and its signature should be
      `func(*loop_vars) =>
      (NDArray or nested List[NDArray] step_output, NDArray or nested List[NDArray] new_loop_vars)`.
  
      `max_iterations` is a scalar that defines the maximum number of iterations allowed.
  
      This function returns two lists.
      The first list has the length of `|step_output|`,
      in which the i-th element are all i-th elements of
      `step_output` from all steps, stacked along axis 0.
      The second list has the length of `|loop_vars|`,
      which represents final states of loop variables.
  
      .. warning::
  
         For now, the axis 0 of all NDArrays in the first list are `max_iterations`,
         due to lack of dynamic shape inference.
  
      .. warning::
  
         When `cond` is never satisfied, we assume `step_output` is empty,
         because it cannot be inferred. This is different from the symbolic version.
  
      Parameters
      ----------
      cond: a Python function.
          The loop condition.
      func: a Python function.
          The loop body.
      loop_vars: an NDArray or nested lists of NDArrays.
          The initial values of the loop variables.
      max_iterations: a python int.
          Maximum number of iterations.
  
      Returns
      ------
      outputs: an NDArray or nested lists of NDArrays
          stacked output from each step
      states: an NDArray or nested lists of NDArrays
          final state
  
      Examples
      --------
      >>> cond = lambda i, s: i <= 5
      >>> func = lambda i, s: ([i + s], [i + 1, s + i])
      >>> loop_vars = (mx.np.array([0], dtype="int64"), mx.np.array([1], dtype="int64"))
      >>> outputs, states = mx.npx.while_loop(cond, func, loop_vars, max_iterations=10)
      >>> outputs
      [array([[ 1],
             [ 2],
             [ 4],
             [ 7],
             [11],
             [16],
             [ 0],
             [ 0],
             [ 0],
             [ 0]], dtype=int64)]
      >>> states
      [array([6], dtype=int64), array([16], dtype=int64)]
      """
  
      def _to_python_scalar(inputs, type_, name):
          """Converts "inputs", possibly typed mxnet NDArray, a numpy ndarray, other python types,
          to the given type
          """
          if isinstance(inputs, np_ndarray):
              inputs = inputs.item()
          try:
              inputs = type_(inputs)
          except:
              raise ValueError(
                  f"Cannot convert {name} to python {type_.__name__}"
              )
          return inputs
  
      def _cond_wrapper(loop_vars):
          if dc.is_deferred_compute():
              result = cond(*loop_vars).astype("int")
          else:
              with ag.pause(), dc.context():
                  result = cond(*loop_vars).astype("int")
          flatten_out, _ = _flatten(result, "while_loop output")
          out = dc.get_symbol(flatten_out)
          dc.clear(flatten_out)
          return [], [out], [], []
  
      def _func_wrapper(loop_vars):
          """This wrapper unifies
               "func: loop_vars -> new_loop_vars"
           and "func: loop_vars -> (step_output, new_loop_vars)"
          into "func: loop_vars -> (None or tuple of step_outputs, tuple of new_loop_vars)
          """
          if dc.is_deferred_compute():
              step_output, new_loop_vars = func(*loop_vars)
          else:
              with ag.pause(), dc.context():
                  step_output, new_loop_vars = func(*loop_vars)
          if step_output is None:
              step_output = []
          if new_loop_vars is None:
              new_loop_vars = []
          if isinstance(step_output, tuple):
              step_output = list(step_output)
          if isinstance(new_loop_vars, tuple):
              new_loop_vars = list(new_loop_vars)
          new_loop_vars = _as_list(new_loop_vars)
          if len(loop_vars) != len(new_loop_vars):
              raise ValueError(
                  "The length of loop_vars should be consistent during the loop"
              )
          step_output_flatten, out_fmt = _flatten(
              step_output, "while output"
          )
          new_loop_vars_flatten, var_fmt = _flatten(
              new_loop_vars, "while loop_vars"
          )
          if isinstance(step_output, list):
              if len(step_output) == 0:
                  step_out = []
              else:
                  step_out = [
                      dc.get_symbol(out) for out in step_output_flatten
                  ]
          else:
              step_output_flatten, out_fmt = _flatten(
                  step_output, "while output"
              )
              step_out = [dc.get_symbol(step_output_flatten)]
          if len(new_loop_vars) == 0:
              new_var = []
          else:
              new_var = [
                  dc.get_symbol(var) for var in new_loop_vars_flatten
              ]
          return step_out, new_var, out_fmt, var_fmt
  
      def _create_subgraph(graph_vars, graph_func, subgraph_name):
          subgraph_name = _get_unique_subgraph_name(subgraph_name)
          with AttrScope(__subgraph_name__=subgraph_name):
              # create new variables with the same name,
              # them feed them to the given func
              flatten_data, data_fmt = _flatten(
                  graph_vars, "foreach input"
              )
              real_data = [
                  ele.copy().detach() if ele is not None else None
                  for ele in flatten_data
              ]
              data_names = [
                  "data_subgraph{}".format(i)
                  for i, ele in enumerate(real_data)
              ]
              symbol_data = [
                  symbol.var(name).as_np_ndarray()
                  for arg, name in zip(real_data, data_names)
              ]
              dc.set_variable(real_data, symbol_data)
              new_graph_vars = _regroup(real_data, data_fmt)
              outputs, final_state, out_fmt, var_fmt = graph_func(
                  new_graph_vars
              )
              # first `num_out_data` elements belong to `outputs`
              # other elements belong to `final_state`
              num_out_data = len(outputs)
              num_outputs = len(outputs) + len(final_state)
              # group all outputs of graph_func
              graph = _construct_subgraph(outputs, final_state)
          return graph, num_out_data, num_outputs, out_fmt, var_fmt
  
      flatten_loop_vars, init_loop_var_fmt = _flatten(
          loop_vars, "while loop_vars"
      )
  
      def _union_inputs(*graphs):
          # Given a list of graphs, each whose inputs are either from loop_vars or other variables.
          # 1) calculate a list `inputs`, the union of their inputs.
          # 2) for each graph, determine in which indices their inputs reside in `inputs`
          # 3) for each variable in the input of `graph`, find which index it is
          inputs = []  # List[Symbol], result of 1)
          locs = (
              []
          )  # List[Tuple(List[Int], List[Int])], a list of tuples,
          # where tuples are results of 2) and 3)
          input_id_to_loc = (
              {}
          )  # Dict[int, int], given id(sym), input_id_to_loc maps it
          # to a `loc`, where inputs[loc] = sym
          for graph in graphs:
              # some loop_vars are inputs to `graph`, some are not
              name_to_loop_vars = {
                  "data_subgraph{}".format(i): ele
                  for i, ele in enumerate(flatten_loop_vars)
              }
              # also we collect the mapping from var's name to var's loc in loop_vars
              name_to_var_locs = {
                  "data_subgraph{}".format(i): i
                  for i, ele in enumerate(flatten_loop_vars)
              }
              # collect arguments for each subgraph
              input_locs = []  # results from the second step
              var_locs = [-1] * len(
                  flatten_loop_vars
              )  # results from the third step
              subg_input_names = graph.list_inputs()
              for name in subg_input_names:
                  assert (
                      name in name_to_loop_vars
                  )  # it should obviously hold
                  array = name_to_loop_vars[name]
                  # do 2), and 1) is implicitly done
                  if id(array) in input_id_to_loc:
                      loc = input_id_to_loc[id(array)]
                  else:
                      loc = len(input_id_to_loc)
                      inputs.append(array)
                      input_id_to_loc[id(array)] = loc
                  input_locs.append(loc)
                  # do 3)
                  if name in name_to_var_locs:
                      var_locs[name_to_var_locs[name]] = (
                          len(input_locs) - 1
                      )
                      name_to_var_locs.pop(name, None)
              locs.append((input_locs, var_locs))
          return inputs, locs
  
      if max_iterations is None:
          raise ValueError("max_iterations should be specified")
      max_iterations = _to_python_scalar(
          max_iterations, int, "max_iteration"
      )
      # It should be work as fine if loop_vars are empty I guess,
      # but it is semantically unnecessary to include this case.
      if isinstance(loop_vars, (list, tuple)):
          if len(loop_vars) == 0:
              raise ValueError(
                  "loop_vars should contain at least one element"
              )
      else:
          assert isinstance(loop_vars, np_ndarray), (
              "loop_vars should be either mxnet.numpy.ndarray"
              " or list/tuple of mxnet.numpy.ndarray"
          )
          loop_vars = [loop_vars]
      # create graph for `cond'
      cond_g, num_out_data, num_outputs, _, _ = _create_subgraph(
          loop_vars, _cond_wrapper, name + "_cond"
      )
      assert num_out_data == 0
      assert num_outputs == 1
      # create graph for `func`
      func_g, num_out_data, num_outputs, out_fmt, _ = _create_subgraph(
          loop_vars, _func_wrapper, name + "_func"
      )
      # find symbols used in either cond_g or func_g
      input_vars, (
          (cond_input_locs, _),
          (func_input_locs, func_var_locs),
      ) = _union_inputs(cond_g, func_g)
      for i_th, loc in enumerate(func_var_locs, 1):
          if loc == -1:
              raise ValueError(
                  f"The {i_th}-th loop_var doesn't involve into the computation"
              )
      result = _api_internal.while_loop(
          cond_g.handle,
          func_g.handle,
          *input_vars,
          max_iterations,
          cond_input_locs,
          func_input_locs,
          func_var_locs,
          num_out_data,
          num_outputs,
      )
      if isinstance(result, np_ndarray):
          ret = [result]
      else:
          ret = list(result)
      outputs = [ret[i] for i in range(num_out_data)]
      outputs = _regroup(outputs, out_fmt)
      final_loop_vars = [
          ret[i] for i in range(num_out_data, num_outputs)
      ]
      final_loop_vars = _regroup(final_loop_vars, init_loop_var_fmt)
      return outputs, final_loop_vars
  
file_path: mxnet/python/mxnet/ndarray/numpy_extension/control_flow.py
imports: from .None import _api_internal
  from .._internal import NDArrayBase
  from ...util import set_module
  from ...numpy import ndarray
  from ...symbol import Symbol
  from ...base import _as_list
  from ...None import symbol, _deferred_compute, autograd
  from ...attribute import AttrScope, current
label_desc: Use image node_id 6 for calling a global function with example usage: while_loop(cond, func, loop_vars, max_iterations, name) and returns: outputs,  final_loop_vars, inputs, list,  list,  list,  list, step_out,  new_var,  out_fmt,  var_fmt, graph,  num_out_data,  num_outputs,  out_fmt,  var_fmt, inputs,  locs