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 Symbol or nested lists of Symbols 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 symbol,
    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) => Symbol`.

    `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) =>
    (Symbol or nested List[Symbol] step_output, Symbol or nested List[Symbol] 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 Symbols in the first list are `max_iterations`,
       due to lack of dynamic shape inference.

    .. warning::

       Even if `cond` is never satisfied,
       while_loop returns a list of outputs with inferred dtype and shape.
       This is different from the Symbol version,
       where in this case `step_outputs` are assumed as an empty list.

    Parameters
    ----------
    cond: a Python function.
        The loop condition.
    func: a Python function.
        The loop body.
    loop_vars: a Symbol or nested lists of Symbol.
        The initial values of the loop variables.
    max_iterations: a python int.
        Maximum number of iterations.

    Returns
    ------
    outputs: a Symbol or nested lists of Symbols
        stacked output from each step
    states: a Symbol or nested lists of Symbols
        final state

    Examples
    --------
    >>> cond = lambda i, s: i <= 5
    >>> func = lambda i, s: ([i + s], [i + 1, s + i])
    >>> loop_vars = (mx.sym.var('i'), mx.sym.var('s'))
    >>> outputs, states = mx.sym.contrib.while_loop(cond, func, loop_vars, max_iterations=10)

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 Symbol or nested lists of Symbols 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 symbol,
    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) => Symbol`.

    `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) =>
    (Symbol or nested List[Symbol] step_output, Symbol or nested List[Symbol] 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 Symbols in the first list are `max_iterations`,
       due to lack of dynamic shape inference.

    .. warning::

       Even if `cond` is never satisfied,
       while_loop returns a list of outputs with inferred dtype and shape.
       This is different from the Symbol version,
       where in this case `step_outputs` are assumed as an empty list.

    Parameters
    ----------
    cond: a Python function.
        The loop condition.
    func: a Python function.
        The loop body.
    loop_vars: a Symbol or nested lists of Symbol.
        The initial values of the loop variables.
    max_iterations: a python int.
        Maximum number of iterations.

    Returns
    ------
    outputs: a Symbol or nested lists of Symbols
        stacked output from each step
    states: a Symbol or nested lists of Symbols
        final state

    Examples
    --------
    >>> cond = lambda i, s: i <= 5
    >>> func = lambda i, s: ([i + s], [i + 1, s + i])
    >>> loop_vars = (mx.sym.var('i'), mx.sym.var('s'))
    >>> outputs, states = mx.sym.contrib.while_loop(cond, func, loop_vars, max_iterations=10)
    """
    def _to_python_scalar(inputs, type_, name):
        """Converts "inputs", possibly typed mxnet NDArray, a numpy ndarray, other python types,
        to the given type
        """
        if hasattr(inputs, "asscalar"):
            inputs = inputs.asscalar()
        try:
            inputs = type_(inputs)
        except:
            raise ValueError("Cannot convert %s to python %s" % (name, type_.__name__))
        return inputs

    def _cond_wrapper(loop_vars):
        result = cond(*loop_vars)
        if not isinstance(result, Symbol):
            raise ValueError("Return of cond must be a Symbol")
        return [], [result], [], []

    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 -> (list of step_outputs, tuple of new_loop_vars)
        """
        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)
        step_output, out_fmt = _flatten(step_output, "while output")
        new_loop_vars, var_fmt = _flatten(new_loop_vars, "while loop_vars")
        if len(loop_vars) != len(new_loop_vars):
            raise ValueError("The number of loop_vars should be consistent during the loop")
        return step_output, new_loop_vars, 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
            graph_vars, var_fmt = _flatten(graph_vars, "while loop_vars")
            new_graph_vars = [symbol.var(_get_sym_uniq_name(sym)) for sym in graph_vars]
            new_graph_vars, _ = _regroup(new_graph_vars, var_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)
            # nnvm cut-graph does not allow inputs and outputs overlap
            # so we calculate the name of inputs, and copy outputs once it overlaps with inputs
            # group all outputs of graph_func
            all_input_names = symbol.Group(outputs + final_state).list_inputs()
            in_input = lambda x: x.name in all_input_names
            in_graph = lambda x: x.list_attr().get("__subgraph_name__", "") == subgraph_name
            make_identity = lambda x: symbol.op.identity(x) if in_input(x) or not in_graph(x) \
                                      else x
            graph = symbol.Group(list(map(make_identity, 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")
    _check_data(flatten_loop_vars, symbol.Symbol,
                "loop_vars should be a symbol or a nested list of symbols")

    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 = {_get_sym_uniq_name(sym): sym for sym in flatten_loop_vars}
            # other inputs to `graph` created by cut_graph
            name_to_cut_g_syms = {sym.list_outputs()[0]: sym for sym in _cut_subgraph(graph)}
            # input_syms: all inputs to the `graph`
            name_to_input_syms = {sym.name: sym for sym in _get_graph_inputs(graph)}
            # also we collect the mapping from var's name to var's loc in loop_vars
            name_to_var_locs = {_get_sym_uniq_name(sym): i for i, sym 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()
            assert len(set(subg_input_names)) == len(subg_input_names), \
                    "The inputs of the subgraph don't have unique names: " + str(subg_input_names)
            for name in subg_input_names:
                assert name in name_to_input_syms   # it should obviously hold
                # name -> sym
                if name in name_to_loop_vars:
                    sym = name_to_loop_vars[name]
                elif name in name_to_cut_g_syms:
                    sym = name_to_cut_g_syms[name]
                else:
                    sym = copy.deepcopy(name_to_input_syms[name])
                # do 2), and 1) is implicitly done
                if id(sym) in input_id_to_loc:
                    loc = input_id_to_loc[id(sym)]
                else:
                    loc = len(input_id_to_loc)
                    inputs.append(sym)
                    input_id_to_loc[id(sym)] = 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
            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 len(loop_vars) == 0:
        raise ValueError("loop_vars should contain at least one element")
    # 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_syms, ((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("The %d-th loop_var doesn't involve into the computation" % i_th)
    result = symbol._internal._while_loop(
        cond_g,
        func_g,
        *input_syms,
        max_iterations=max_iterations,
        cond_input_locs=cond_input_locs,
        func_input_locs=func_input_locs,
        func_var_locs=func_var_locs,
        num_out_data=num_out_data,
        num_outputs=num_outputs
    )
    outputs = [result[i] for i in range(num_out_data)]
    outputs, _ = _regroup(outputs, out_fmt)
    final_loop_vars = [result[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