Run an if-then-else using user-defined condition and computation

    This operator simulates a if-like branch which chooses to do one of
    the two customized computations according to the specified condition.

    `pred` is a scalar MXNet Symbol,
    indicating which branch of computation should be used.

    `then_func` is a user-defined function, used as computation of the then branch.
    It produces `outputs`, which is a list of Symbols.
    The signature of `then_func` should be
    `then_func() => nested List[Symbol]`.

    `else_func` is a user-defined function, used as computation of the else branch.
    It produces `outputs`, which is a list of Symbols.
    The signature of `else_func` should be
    `else_func() => nested List[Symbol]`.

    The `outputs` produces by `then_func` and `else_func` should have the same number
    of elements, all of which should be in the same shape, of the same dtype and stype.

    This function returns a list of symbols, representing the computation result.

    Parameters
    ----------
    pred: a MXNet Symbol representing a scalar.
        The branch condition.
    then_func: a Python function.
        The computation to be executed if `pred` is true.
    else_func: a Python function.
        The computation to be executed if `pred` is false.

    Returns
    -------
    outputs: a Symbol or nested lists of Symbols, representing the result of computation.

    Examples
    --------
    >>> a, b = mx.sym.var('a'), mx.sym.var('b')
    >>> pred = a * b < 5
    >>> then_func = lambda: (a + 5) * (b + 5)
    >>> else_func = lambda: (a - 5) * (b - 5)
    >>> outputs = mx.sym.contrib.cond(pred, then_func, else_func)

def cond(pred, then_func, else_func, name="cond"):
    """Run an if-then-else using user-defined condition and computation

    This operator simulates a if-like branch which chooses to do one of
    the two customized computations according to the specified condition.

    `pred` is a scalar MXNet Symbol,
    indicating which branch of computation should be used.

    `then_func` is a user-defined function, used as computation of the then branch.
    It produces `outputs`, which is a list of Symbols.
    The signature of `then_func` should be
    `then_func() => nested List[Symbol]`.

    `else_func` is a user-defined function, used as computation of the else branch.
    It produces `outputs`, which is a list of Symbols.
    The signature of `else_func` should be
    `else_func() => nested List[Symbol]`.

    The `outputs` produces by `then_func` and `else_func` should have the same number
    of elements, all of which should be in the same shape, of the same dtype and stype.

    This function returns a list of symbols, representing the computation result.

    Parameters
    ----------
    pred: a MXNet Symbol representing a scalar.
        The branch condition.
    then_func: a Python function.
        The computation to be executed if `pred` is true.
    else_func: a Python function.
        The computation to be executed if `pred` is false.

    Returns
    -------
    outputs: a Symbol or nested lists of Symbols, representing the result of computation.

    Examples
    --------
    >>> a, b = mx.sym.var('a'), mx.sym.var('b')
    >>> pred = a * b < 5
    >>> then_func = lambda: (a + 5) * (b + 5)
    >>> else_func = lambda: (a - 5) * (b - 5)
    >>> outputs = mx.sym.contrib.cond(pred, then_func, else_func)
    """

    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
            new_graph_vars = [symbol.var(sym.name) for sym in graph_vars]
            outputs = graph_func(*new_graph_vars)
            outputs, out_fmt = _flatten(outputs, "cond outputs")
            num_outputs = len(outputs)
            # 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).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)))
        return graph, num_outputs, out_fmt

    def _union_inputs(*graphs):
        # Given a list of graphs, each whose inputs are either from input_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 input_vars are inputs to `graph`, some are not
            name_to_input_vars = {sym.name: sym for sym in inputs}
            # 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)}
            # collect arguments for each subgraph
            input_locs = []                         # results from the second step
            for name in graph.list_inputs():
                assert name in name_to_input_syms   # it should obviously hold
                # name -> sym
                if name in name_to_input_vars:
                    sym = name_to_input_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)
            locs.append(input_locs)
        return inputs, locs
    inputs = []
    # create graph for `cond_func'
    cond_g, cond_num_outputs, _ = _create_subgraph(inputs, lambda: pred, name + "_pred")
    if cond_num_outputs != 1:
        raise ValueError("pred should always be a single output")
    # create graph for `then`
    then_g, then_num_outputs, then_fmt = _create_subgraph(inputs, then_func, name + "_then")
    # create graph for `else`
    else_g, else_num_outputs, _ = _create_subgraph(inputs, else_func, name + "_else")
    if then_num_outputs != else_num_outputs:
        raise ValueError("Number of outputs differs between then-branch and else-branch")
    # find symbols used in either cond_g or func_g
    input_syms, (cond_input_locs, then_input_locs, else_input_locs) = \
        _union_inputs(cond_g, then_g, else_g)
    result = symbol._internal._cond(
        # [cond, then_g, else_g, *input_syms]
        cond_g,
        then_g,
        else_g,
        *input_syms,
        cond_input_locs=cond_input_locs,
        then_input_locs=then_input_locs,
        else_input_locs=else_input_locs,
        num_outputs=then_num_outputs
    )
    outputs = [result[i] for i in range(then_num_outputs)]
    outputs, _ = _regroup(outputs, then_fmt)
    return outputs