parameters: (pred, then_func, else_func, inputs, name="cond")
return_statement: return outputs
docstring: 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
docstring_summary: Run an if-then-else using user-defined condition and computation
docstring_tokens: Run an if - then - else using user - defined condition and computation
function: 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
function_tokens: def cond ( pred , then_func , else_func , inputs , name = "cond" ) : 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