Run a for loop with user-defined computation over Symbols on dimension 0.

    This operator simulates a for loop and body has the computation for an iteration
    of the for loop. It runs the computation in body on each slice from the input
    NDArrays.

    body takes two arguments as input and outputs a tuple of two elements,
    as illustrated below:

    out, states = body(data1, states)

    data1 can be either a symbol or a list of symbols. If data is a symbol,
    data1 is a symbol. Otherwise, data1 is a list of symbols and has the same
    size as data. states is a list of symbols and have the same size as init_states.
    Similarly, out can be either a symbol or a list of symbols, which are concatenated
    as the first output of foreach; states from the last execution of body
    are the second output of foreach.

    foreach can output only output data or states. If a user only wants states,
    the body function can return ([], states). Similarly, if a user only wants
    output data, the body function can return (out, []).

    The computation done by this operator is equivalent to the pseudo code below
    when the input data is NDArray::

        states = init_states
        outs = []
        for i in data.shape[0]:
            s = data[i]
            out, states = body(s, states)
            outs.append(out)
        outs = stack(*outs)


    Parameters
    ----------
    body : a Python function.
        Define computation in an iteration.
    data: a symbol or a list of symbols.
        The input data.
    init_states: a Symbol or nested lists of symbols.
        The initial values of the loop states.
    name: string.
        The name of the operator.

    Returns
    -------
    outputs: a Symbol or nested lists of Symbols.
        The output data concatenated from the output of all iterations.
    states: a Symbol or nested lists of Symbols.
        The loop states in the last iteration.

    Examples
    --------
    >>> step = lambda data, states: (data + states[0], [states[0] * 2])
    >>> data = mx.sym.var('data')
    >>> states = [mx.sym.var('state')]
    >>> outs, states = mx.sym.contrib.foreach(step, data, states)

def foreach(body, data, init_states, name="foreach"):
    """Run a for loop with user-defined computation over Symbols on dimension 0.

    This operator simulates a for loop and body has the computation for an iteration
    of the for loop. It runs the computation in body on each slice from the input
    NDArrays.

    body takes two arguments as input and outputs a tuple of two elements,
    as illustrated below:

    out, states = body(data1, states)

    data1 can be either a symbol or a list of symbols. If data is a symbol,
    data1 is a symbol. Otherwise, data1 is a list of symbols and has the same
    size as data. states is a list of symbols and have the same size as init_states.
    Similarly, out can be either a symbol or a list of symbols, which are concatenated
    as the first output of foreach; states from the last execution of body
    are the second output of foreach.

    foreach can output only output data or states. If a user only wants states,
    the body function can return ([], states). Similarly, if a user only wants
    output data, the body function can return (out, []).

    The computation done by this operator is equivalent to the pseudo code below
    when the input data is NDArray::

        states = init_states
        outs = []
        for i in data.shape[0]:
            s = data[i]
            out, states = body(s, states)
            outs.append(out)
        outs = stack(*outs)


    Parameters
    ----------
    body : a Python function.
        Define computation in an iteration.
    data: a symbol or a list of symbols.
        The input data.
    init_states: a Symbol or nested lists of symbols.
        The initial values of the loop states.
    name: string.
        The name of the operator.

    Returns
    -------
    outputs: a Symbol or nested lists of Symbols.
        The output data concatenated from the output of all iterations.
    states: a Symbol or nested lists of Symbols.
        The loop states in the last iteration.

    Examples
    --------
    >>> step = lambda data, states: (data + states[0], [states[0] * 2])
    >>> data = mx.sym.var('data')
    >>> states = [mx.sym.var('state')]
    >>> outs, states = mx.sym.contrib.foreach(step, data, states)
    """

    flatten_data, data_fmt = _flatten(data, "foreach input")
    _check_data(flatten_data, symbol.Symbol,
                "data should be a symbol or a nested list of symbols")
    init_flatten_states, init_state_fmt = _flatten(init_states, "foreach states")
    _check_data(init_flatten_states, symbol.Symbol,
                "init_states should be a symbol or a nested list of symbols")

    # If the input python function references to the symbols outside
    # the python function, we need to prune the computation graph constructed from
    # the function. One way of doing it is to mark the nodes in the computation graph
    # with AttrScope and prune the nodes without the special attribute.
    name = _get_unique_subgraph_name(name)
    with AttrScope(__subgraph_name__=name):
        in_eles = [symbol.var(_get_sym_uniq_name(sym)) for sym in flatten_data]
        in_eles, _ = _regroup(in_eles, data_fmt)
        states = [symbol.var(_get_sym_uniq_name(s)) for s in init_flatten_states]
        states, _ = _regroup(states, copy.deepcopy(init_state_fmt))
        sym_out, sym_states = body(in_eles, states)

        sym_out, out_fmt = _flatten(sym_out, "foreach output")
        sym_states, state_fmt = _flatten(sym_states, "foreach loop_vars")
        assert init_state_fmt == state_fmt, "The input and output loop_vars have different format"
        _check_data(sym_out, symbol.Symbol,
                    "the output should be an NDArray or a nested list of NDArrays")
        _check_data(sym_states, symbol.Symbol,
                    "the output states should be an NDArray or a nested list of NDArrays")
        num_out_data = len(sym_out)
        num_states = len(sym_states)
        num_outputs = num_out_data + num_states
        g = _construct_subgraph(sym_out, sym_states, name)

    input_syms = _get_graph_inputs(g)
    cut_syms = _cut_subgraph(g)
    input_syms = _get_graph_inputs(g)

    # Here we need to find out how the input symbols are ordered as well as
    # where the loop states are located in the list of inputs.

    # This dict contains the symbols of the subgraph.
    input_syms = {sym.name:sym for sym in input_syms}
    gin_names = input_syms.keys()
    # This array contains the symbols for the inputs of foreach.
    # They are ordered according to the inputs of the subgraph.
    state_names = [_get_sym_uniq_name(sym) for sym in init_flatten_states]
    data_names = [_get_sym_uniq_name(sym) for sym in flatten_data]
    cut_var_map = {sym.list_outputs()[0]:sym for sym in cut_syms}
    cut_var_names = cut_var_map.keys()

    subg_input_names = g.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)
    # ordered_ins contains input symbols in the following order:
    # data_syms, state_syms, followed by cut_vars and vars in the closure.
    ordered_ins = [x for x in flatten_data]
    # this defines the location of data_syms in the list of subgraph inputs
    in_data_locs = []
    for dname in data_names:
        # Some data may not be used.
        if dname in subg_input_names:
            in_data_locs.append(subg_input_names.index(dname))
        else:
            raise AssertionError("the data arrays have to be used in the loop body")

    ordered_ins.extend(init_flatten_states)
    # this defines the location of state_syms in the list of subgraph inputs.
    in_state_locs = []
    for sname in state_names:
        # Some state may not be used.
        if sname in subg_input_names:
            in_state_locs.append(subg_input_names.index(sname))
        else:
            raise AssertionError("the state arrays have to be used in the loop body")

    remain_locs = []
    for in_name in subg_input_names:
        assert in_name in gin_names, "The input variable %s can't be found in graph inputs: %s" \
                % (in_name, str(gin_names))
        if in_name in cut_var_names:
            ordered_ins.append(cut_var_map[in_name])
            remain_locs.append(subg_input_names.index(in_name))
        elif in_name not in data_names and in_name not in state_names:
            # The remaining inputs are the variable nodes created inside the UDF.
            # The subgraph can't have nodes shared with the main graph. As such,
            # we need to make a copy of these variable nodes.
            assert in_name in gin_names
            ordered_ins.append(copy.deepcopy(input_syms[in_name]))
            remain_locs.append(subg_input_names.index(in_name))

    ret = symbol._internal._foreach(g, *ordered_ins, num_outputs=num_outputs,
                                    num_out_data=num_out_data, in_state_locs=in_state_locs,
                                    in_data_locs=in_data_locs, remain_locs=remain_locs)
    outs = []
    for i in range(num_outputs - num_states):
        outs.append(ret[i])
    outs, _ = _regroup(outs, out_fmt)
    states = []
    for i in range(num_states):
        states.append(ret[num_outputs - num_states + i])
    states, _ = _regroup(states, state_fmt)

    return (outs, states)