Source code for mxnet.symbol.contrib

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# coding: utf-8
# pylint: disable=wildcard-import, unused-wildcard-import,redefined-outer-name
"""Contrib Symbol API of MXNet."""
import math
import ctypes
import copy

from .random import uniform
from .symbol import Symbol
try:
    from .gen_contrib import *
except ImportError:
    pass

from . import symbol
from ..base import _LIB, check_call
from ..base import SymbolHandle, _as_list
from ..attribute import AttrScope

__all__ = ["rand_zipfian", "foreach", "while_loop", "cond"]

[docs]def rand_zipfian(true_classes, num_sampled, range_max): """Draw random samples from an approximately log-uniform or Zipfian distribution. This operation randomly samples *num_sampled* candidates the range of integers [0, range_max). The elements of sampled_candidates are drawn with replacement from the base distribution. The base distribution for this operator is an approximately log-uniform or Zipfian distribution: P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1) This sampler is useful when the true classes approximately follow such a distribution. For example, if the classes represent words in a lexicon sorted in decreasing order of \ frequency. If your classes are not ordered by decreasing frequency, do not use this op. Additionaly, it also returns the number of times each of the \ true classes and the sampled classes is expected to occur. Parameters ---------- true_classes : Symbol The target classes in 1-D. num_sampled: int The number of classes to randomly sample. range_max: int The number of possible classes. Returns ------- samples: Symbol The sampled candidate classes in 1-D `int64` dtype. expected_count_true: Symbol The expected count for true classes in 1-D `float64` dtype. expected_count_sample: Symbol The expected count for sampled candidates in 1-D `float64` dtype. Examples -------- >>> true_cls = mx.nd.array([3]) >>> samples, exp_count_true, exp_count_sample = mx.nd.contrib.rand_zipfian(true_cls, 4, 5) >>> samples [1 3 3 3] >>> exp_count_true [ 0.12453879] >>> exp_count_sample [ 0.22629439 0.12453879 0.12453879 0.12453879] """ assert(isinstance(true_classes, Symbol)), "unexpected type %s" % type(true_classes) log_range = math.log(range_max + 1) rand = uniform(0, log_range, shape=(num_sampled,), dtype='float64') # make sure sampled_classes are in the range of [0, range_max) sampled_classes = (rand.exp() - 1).astype('int64') % range_max true_classes = true_classes.astype('float64') expected_prob_true = ((true_classes + 2.0) / (true_classes + 1.0)).log() / log_range expected_count_true = expected_prob_true * num_sampled # cast sampled classes to fp64 to avoid interget division sampled_cls_fp64 = sampled_classes.astype('float64') expected_prob_sampled = ((sampled_cls_fp64 + 2.0) / (sampled_cls_fp64 + 1.0)).log() / log_range expected_count_sampled = expected_prob_sampled * num_sampled return sampled_classes, expected_count_true, expected_count_sampled
def _flatten(args, inout_str): if isinstance(args, symbol.Symbol): length = len(args.list_outputs()) length = length if length > 1 else 0 return [args], int(length) assert isinstance(args, (list, tuple)), \ "%s must be (nested) list of Symbol, " \ "but got %s of type %s"%(inout_str, str(args), str(type(args))) flat = [] fmts = [] for i in args: arg, fmt = _flatten(i, inout_str) flat.extend(arg) fmts.append(fmt) return flat, fmts def _regroup(args, fmt): if isinstance(fmt, int): if fmt == 0: return args[0], args[1:] return args[:fmt], args[fmt:] assert isinstance(args, (list, tuple)), \ "output must be (nested) list of Symbol, " \ "but got %s of type %s"%(str(args), str(type(args))) ret = [] for i in fmt: res, args = _regroup(args, i) ret.append(res) return ret, args # We want to generate a unique name for input symbols to a control flow # operator. The names are generated on purpose differently from the symbols # cut from the graph. def _get_sym_uniq_name(sym): return '{}-{}'.format(sym.name, sym.attr('_value_index')) def _get_graph_inputs(subg): num_handles = ctypes.c_int(0) handles = ctypes.POINTER(SymbolHandle)() check_call(_LIB.MXSymbolGetInputSymbols(subg.handle, ctypes.byref(handles), ctypes.byref(num_handles))) syms = [] for i in range(num_handles.value): s = Symbol(ctypes.cast(handles[i], SymbolHandle)) syms.append(s) return syms def _cut_subgraph(subg): num_handles = ctypes.c_int(0) handles = ctypes.POINTER(SymbolHandle)() check_call(_LIB.MXSymbolCutSubgraph(subg.handle, ctypes.byref(handles), ctypes.byref(num_handles))) syms = [] for i in range(num_handles.value): s = Symbol(ctypes.cast(handles[i], SymbolHandle)) syms.append(s) return syms def _get_unique_subgraph_name(subgraph_name): attrs = AttrScope._current.value._attr if attrs.get("__subgraph_name__", "") != "": subgraph_name = "".join([attrs["__subgraph_name__"], "$", subgraph_name]) AttrScope._subgraph_names[subgraph_name] += 1 subgraph_name = subgraph_name + str(AttrScope._subgraph_names[subgraph_name] - 1) return subgraph_name # This construct a subgraph for given output nodes. # If an output node is one of the input nodes, we call identity to make sure # that outputs nodes are different from input nodes. def _construct_subgraph(sym_out, sym_states, name): sym_out = _as_list(sym_out) sym_states = _as_list(sym_states) all_outputs = [] all_outputs.extend(sym_out) all_outputs.extend(sym_states) g = symbol.Group(all_outputs) flat_out = [] all_input_names = g.list_inputs() output_names = {o.name for o in sym_out} for o in sym_out: if o.name in all_input_names or o.list_attr().get("__subgraph_name__", "") != name: flat_out.append(symbol.op.identity(o)) else: flat_out.append(o) for s in sym_states: if s.name in all_input_names or s.name in output_names or \ s.list_attr().get("__subgraph_name__", "") != name: flat_out.append(symbol.op.identity(s)) else: flat_out.append(s) return symbol.Group(flat_out) def _check_data(inputs, in_type, msg): is_NDArray_or_list = True if isinstance(inputs, list): for i in inputs: if not isinstance(i, in_type): is_NDArray_or_list = False break else: is_NDArray_or_list = isinstance(inputs, in_type) assert is_NDArray_or_list, msg
[docs]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)
[docs]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
[docs]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