code: def cond(pred, then_func, else_func, inputs, 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 NDArray,
      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 NDArrays.
      The signature of `then_func` should be
      `then_func() => NDArray or nested List[NDArray]`.
  
      `else_func` is a user-defined function, used as computation of the else branch.
      It produces `outputs`, which is a list of NDArrays.
      The signature of `else_func` should be
      `else_func() => NDArray or nested List[NDArray]`.
  
      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 Python function.
          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: an NDArray or nested lists of NDArrays, representing the result of computation.
  
      Examples
      --------
      >>> a, b = mx.np.array([1]), mx.np.array([2])
      >>> pred = a * b < 5
      >>> then_func = lambda: (a + 5) * (b + 5)
      >>> else_func = lambda: (a - 5) * (b - 5)
      >>> outputs = mx.npx.cond(pred, then_func, else_func)
      >>> outputs[0]
      42.0
      """
  
      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, "cond 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)
              if dc.is_deferred_compute():
                  outputs = graph_func(*new_graph_vars)
                  if "pred" in subgraph_name:
                      outputs = outputs.astype("int")
              else:
                  with ag.pause(), dc.context():
                      outputs = graph_func(*new_graph_vars)
                      if "pred" in subgraph_name:
                          outputs = outputs.astype("int")
              outputs, out_fmt = _flatten(outputs, "cond outputs")
              num_outputs = len(outputs)
              sym_out = [
                  dc.get_symbol(out_data) for out_data in outputs
              ]
              dc.clear(outputs)
              graph = _construct_subgraph(sym_out, [])
          return graph, num_outputs, out_fmt
  
      flatten_inputs, _ = _flatten(inputs, "while loop_vars")
  
      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_syms = {
                  "data_subgraph{}".format(i): ele
                  for i, ele in enumerate(flatten_inputs)
              }
              # 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
                  array = name_to_input_syms[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)
              locs.append(input_locs)
          return inputs, locs
  
      if isinstance(inputs, (list, tuple)):
          if len(inputs) == 0:
              raise ValueError(
                  "inputs should contain at least one element"
              )
      else:
          assert isinstance(inputs, np_ndarray), (
              "inputs should be either mxnet.numpy.ndarray"
              " or list/tuple of mxnet.numpy.ndarray"
          )
          inputs = [inputs]
      # create graph for `cond_func'
      cond_g, cond_num_outputs, _ = _create_subgraph(
          inputs, 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
      union_inputs, (
          cond_input_locs,
          then_input_locs,
          else_input_locs,
      ) = _union_inputs(cond_g, then_g, else_g)
      result = _api_internal.cond(
          cond_g.handle,
          then_g.handle,
          else_g.handle,
          *union_inputs,
          cond_input_locs,
          then_input_locs,
          else_input_locs,
          then_num_outputs,
      )
      if isinstance(result, np_ndarray):
          ret = [result]
      else:
          ret = list(result)
      outputs = [ret[i] for i in range(then_num_outputs)]
      outputs = _regroup(outputs, then_fmt)
      return outputs
  
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 7 for calling a global function with example usage: cond(pred, then_func, else_func, inputs, name) and returns: outputs, graph,  num_outputs,  out_fmt, inputs,  locs