### Step 4. Compile the dummy function ## # We can now compile a dummy function just to see what shared variable # we have and what are their update rules (note that the user has # the option not to pass the shared variable to scan, so we need to # pick them manually and add them to scan) # make the compilation as fast as possible by not applying any # optimization or conversion to C [ note this region is not important # for performance so we can do stuff as unoptimal as we wish ] # extract still missing inputs (there still might be so) and add them # as non sequences at the end of our args fake_nonseqs = [x.type() for x in non_seqs] fake_outputs = scan_utils.clone(outputs, replace=dict(zip(non_seqs, fake_nonseqs))) all_inputs = itertools.ifilter( lambda x: (isinstance(x, gof.Variable) and not isinstance(x, SharedVariable) and not isinstance(x, gof.Constant)), gof.graph.inputs(fake_outputs)) extra_inputs = filter(lambda x: x not in args + fake_nonseqs, all_inputs) non_seqs += extra_inputs ## Note we do not use all_inputs directly since the order of variables ## in args is quite important dummy_args += extra_inputs dummy_outs = outputs if condition is not None:
def remove_constants_and_unused_inputs_scan(node): """ Move constants into the inner graph, and remove unused inputs. Constants that are in the outer graph are represented by a free symbolic variable in the inner graph. If we move them into the inner graph, constant-folding can happen in the inner graph. This is applied only on sequences and non-sequences, not on initial states. """ if not isinstance(node.op, scan_op.Scan): return False op = node.op # We only need to take care of sequences and other arguments st = op.n_seqs st += int(numpy.sum([len(x) for x in op.tap_array[: (op.n_mit_mot + op.n_mit_sot)]])) st += op.n_sit_sot st += op.n_shared_outs op_ins, op_outs = scan_utils.reconstruct_graph(op.inputs, op.outputs) # Corresponds to the initial states, which should stay untouched. # We put those variables aside, and put them back at the end. out_stuff_inner = op_ins[op.n_seqs : st] non_seqs = op_ins[st:] st = op.n_seqs + op.n_mit_mot + op.n_mit_sot + op.n_sit_sot + op.n_nit_sot + op.n_shared_outs + 1 outer_non_seqs = node.inputs[st:] out_stuff_outer = node.inputs[1 + op.n_seqs : st] # To replace constants in the outer graph by clones in the inner graph givens = {} # All the inputs of the inner graph of the new scan nw_inner = [] # Same for the outer graph, initialized w/ number of steps nw_outer = [node.inputs[0]] all_ins = gof.graph.inputs(op_outs) for idx in xrange(op.n_seqs): if ( isinstance(node.inputs[idx + 1], tensor.TensorConstant) and node.inputs[idx + 1].tag.unique_value is not None ): try: # This works if input is a constant that has all entries # equal val = tensor.get_constant_value(node.inputs[idx + 1]) givens[op_ins[idx]] = node.inputs[idx + 1].clone()[0] except TypeError: pass elif op_ins[idx] in all_ins: nw_inner += [op_ins[idx]] nw_outer += [node.inputs[idx + 1]] nw_n_seqs = len(nw_inner) # Add outputs stuff nw_inner += out_stuff_inner nw_outer += out_stuff_outer # Look through non sequences for nw_in, nw_out in zip(non_seqs, outer_non_seqs): if isinstance(nw_out, tensor.Constant): givens[nw_in] = nw_out.clone() elif nw_in in all_ins: nw_inner += [nw_in] nw_outer += [nw_out] if len(nw_inner) != len(op_ins): op_outs = scan_utils.clone(op_outs, replace=givens) nw_info = op.info.copy() nw_info["n_seqs"] = nw_n_seqs # DEBUG CHECK nwScan = scan_op.Scan(nw_inner, op_outs, nw_info) nw_outs = nwScan.make_node(*nw_outer).outputs return nw_outs else: return False
def scan_merge_inouts(node): if not isinstance(node.op, scan_op.Scan): return False a = scan_args(node.inputs, node.outputs, node.op.inputs, node.op.outputs, node.op.info) inp_equiv = {} if has_duplicates(a.outer_in_seqs): new_outer_seqs = [] new_inner_seqs = [] for out_seq, in_seq in zip(a.outer_in_seqs, a.inner_in_seqs): if out_seq in new_outer_seqs: i = new_outer_seqs.index(out_seq) inp_equiv[in_seq] = new_inner_seqs[i] else: new_outer_seqs.append(out_seq) new_inner_seqs.append(in_seq) a.outer_in_seqs = new_outer_seqs a.inner_in_seqs = new_inner_seqs if has_duplicates(a.outer_in_non_seqs): new_outer_nseqs = [] new_inner_nseqs = [] for out_nseq, in_nseq in zip(a.outer_in_non_seqs, a.inner_in_non_seqs): if out_nseq in new_outer_nseqs: i = new_outer_nseqs.index(out_nseq) inp_equiv[in_nseq] = new_inner_nseqs[i] else: new_outer_nseqs.append(out_nseq) new_inner_nseqs.append(in_nseq) a.outer_in_non_seqs = new_outer_nseqs a.inner_in_non_seqs = new_inner_nseqs if len(inp_equiv) > 0: # do the replacement now. The rest will be left to ScanSaveMem inner_inputs = a.inner_inputs outer_inputs = a.outer_inputs info = a.info if info["as_while"]: a_inner_outs = a.inner_outputs + a.cond else: a_inner_outs = a.inner_outputs inner_outputs = scan_utils.clone(a_inner_outs, replace=inp_equiv) orig_outputs = a.outer_outputs op = scan_op.Scan(inner_inputs, inner_outputs, info) outputs = op(*outer_inputs) if not isinstance(outputs, (list, tuple)): outputs = [outputs] na = scan_args(outer_inputs, outputs, op.inputs, op.outputs, op.info) else: na = a # start again left = [] right = [] if has_duplicates(na.outer_in_shared): _left, _right = make_equiv(na.outer_in_shared, na.inner_in_shared) left += _left right += _right if has_duplicates(na.outer_in_sit_sot): _left, _right = make_equiv(na.outer_in_sit_sot, na.inner_in_sit_sot) left += _left right += _right if has_duplicates(na.outer_in_mit_mot): seen = {} for omm, imm, _sl in zip(na.outer_in_mit_mot, na.inner_in_mit_mot, na.mit_mot_in_slices): sl = tuple(_sl) if (omm, sl) in seen: simm = seen[(omm, sl)] left += imm right += simm else: seen[(omm, sl)] = imm if has_duplicates(na.outer_in_mit_sot): seen = {} for oms, ims, _sl in zip(na.outer_in_mit_sot, na.inner_in_mit_sot, na.mit_sot_in_slices): sl = tuple(_sl) if (oms, sl) in seen: sims = seen[(oms, sl)] left += ims right += sims else: seen[(oms, sl)] = ims def map_out(i, o, seen): for si, so in seen: if equal_computations([i], [si], left, right): return so seen.append((i, o)) return o seen = [] na.outer_out_nit_sot = [map_out(i, o, seen) for i, o in zip(na.inner_out_nit_sot, na.outer_out_nit_sot)] seen = [] na.outer_out_sit_sot = [map_out(i, o, seen) for i, o in zip(na.inner_out_sit_sot, na.outer_out_sit_sot)] seen = [] na.outer_out_mit_sot = [map_out(i, o, seen) for i, o in zip(na.inner_out_mit_sot, na.outer_out_mit_sot)] seen = [] new_outer_out_mit_mot = [] for imm, omm, osl in zip(na.inner_out_mit_mot, na.outer_out_mit_mot, na.mit_mot_out_slices): for simm, somm, sosl in seen: if osl == sosl and equal_computations(imm, simm, left, right): new_outer_out_mit_mot.append(somm) break else: seen.append((imm, omm, osl)) new_outer_out_mit_mot.append(omm) na.outer_out_mit_mot = new_outer_out_mit_mot return na.outer_outputs
def remove_constants_and_unused_inputs_scan(node): ''' Move constants into the inner graph, and remove unused inputs. Constants that are in the outer graph are represented by a free symbolic variable in the inner graph. If we move them into the inner graph, constant-folding can happen in the inner graph. This is applied only on sequences and non-sequences, not on initial states. ''' if not isinstance(node.op, scan_op.Scan): return False op = node.op # We only need to take care of sequences and other arguments st = op.n_seqs st += int(numpy.sum([len(x) for x in op.tap_array[:(op.n_mit_mot + op.n_mit_sot)]])) st += op.n_sit_sot st += op.n_shared_outs op_ins, op_outs = scan_utils.reconstruct_graph(op.inputs, op.outputs) # Corresponds to the initial states, which should stay untouched. # We put those variables aside, and put them back at the end. out_stuff_inner = op_ins[op.n_seqs:st] non_seqs = op_ins[st:] st = (op.n_seqs + op.n_mit_mot + op.n_mit_sot + op.n_sit_sot + op.n_nit_sot + op.n_shared_outs + 1) outer_non_seqs = node.inputs[st:] out_stuff_outer = node.inputs[1 + op.n_seqs:st] # To replace constants in the outer graph by clones in the inner graph givens = {} # All the inputs of the inner graph of the new scan nw_inner = [] # Same for the outer graph, initialized w/ number of steps nw_outer = [node.inputs[0]] all_ins = gof.graph.inputs(op_outs) for idx in xrange(op.n_seqs): if (isinstance(node.inputs[idx + 1], tensor.TensorConstant) and node.inputs[idx + 1].tag.unique_value is not None): try: # This works if input is a constant that has all entries # equal givens[op_ins[idx]] = node.inputs[idx + 1].clone()[0] except TypeError: pass elif op_ins[idx] in all_ins: nw_inner += [op_ins[idx]] nw_outer += [node.inputs[idx + 1]] nw_n_seqs = len(nw_inner) # Add outputs stuff nw_inner += out_stuff_inner nw_outer += out_stuff_outer # Look through non sequences for nw_in, nw_out in zip(non_seqs, outer_non_seqs): if isinstance(nw_out, tensor.Constant): givens[nw_in] = nw_out.clone() elif nw_in in all_ins: nw_inner += [nw_in] nw_outer += [nw_out] if len(nw_inner) != len(op_ins): op_outs = scan_utils.clone(op_outs, replace=givens) nw_info = copy.deepcopy(op.info) nw_info['n_seqs'] = nw_n_seqs # DEBUG CHECK nwScan = scan_op.Scan(nw_inner, op_outs, nw_info) nw_outs = nwScan.make_node(*nw_outer).outputs return nw_outs else: return False
def process_node(self, env, node): # this flag tells if there was any change during the last iterations changed = True clean_inputs, clean_outputs = scan_utils.reconstruct_graph(node.op.inputs, node.op.outputs) local_env = gof.Env(clean_inputs, clean_outputs) max_iterations = 2 * len(local_env.toposort()) + 3 counts = 0 to_remove = [] to_replace = [] replace_with_in = [] replace_with_out = [] op = node.op # Construct the list of non_sequences to simplify a few things st = op.n_seqs st += int(numpy.sum([len(x) for x in op.tap_array[: (op.n_mit_mot + op.n_mit_sot)]])) st += op.n_sit_sot st += op.n_shared_outs non_seqs = clean_inputs[st:] st = op.n_seqs + op.n_mit_mot + op.n_mit_sot + op.n_sit_sot + op.n_nit_sot + op.n_shared_outs + 1 outer_non_seqs = node.inputs[st:] assert len(non_seqs) == len(outer_non_seqs) while changed and counts < max_iterations: counts += 1 changed = False for nd in local_env.toposort(): if ( numpy.all( [(x in non_seqs) or (x.owner in to_remove) or isinstance(x, tensor.Constant) for x in nd.inputs] ) and # we can do this because the assumption is that a # viewOp or deepCopyOp will be just at the end of the # function and not somewhere in the middle .. not isinstance(nd.op, theano.compile.ViewOp) and not isinstance(nd.op, theano.compile.DeepCopyOp) and # and we didn't already looked at this node not nd in to_remove ): # We have a candidate node to removable # Step 1. Reconstruct it on outside to_remove.append(nd) outside_ins = [] for x in nd.inputs: if x in non_seqs: outside_ins += [outer_non_seqs[non_seqs.index(x)]] elif x in to_replace: outside_ins += [replace_with_out[to_replace.index(x)]] elif isinstance(x, theano.Constant): outside_ins += [x.clone()] else: raise Exception( ( "Error in the `scan_pushout_non_seq_" "operations`. The optimization tries " "to move some computation fron scan " "which is not allowed to move. Report " "this on theano-users list" ), x, ) nw_outer_node = nd.op.make_node(*outside_ins) # Step 2. Create variables for replacements for idx, y in enumerate(nd.outputs): y_place_holder = scan_utils.safe_new(y, "_replace") to_replace += [y] replace_with_in += [y_place_holder] assert type(y) == type(nw_outer_node.outputs[idx]) replace_with_out += [nw_outer_node.outputs[idx]] changed = True if counts >= max_iterations: raise Exception( "Error in the `scan_pushout_non_seq_operations`." " The optimization exhausted the maximal number " "of iterations allowed!" ) # We need to check all candidate replacements and choose those that # make sense for us # Step 1. which elements of `to_replace` are used by remaining # components of the inner function clean_to_replace = [] clean_replace_with_in = [] clean_replace_with_out = [] existent_nodes = [nd for nd in local_env.toposort() if nd not in to_remove] to_keep = [] for nd in existent_nodes: to_keep += nd.inputs for idx, out in enumerate(to_replace): if out in to_keep and out.owner not in existent_nodes: clean_to_replace += [out] clean_replace_with_in += [replace_with_in[idx]] clean_replace_with_out += [replace_with_out[idx]] if len(clean_to_replace) > 0: # We can finally put an end to all this madness givens = {} nw_outer = [] nw_inner = [] for to_repl, repl_in, repl_out in zip(clean_to_replace, clean_replace_with_in, clean_replace_with_out): if isinstance(repl_out, theano.Constant): repl_in = repl_out.clone() else: nw_inner += [repl_in] nw_outer += [repl_out] givens[to_repl] = repl_in _op_outs = scan_utils.clone(clean_outputs, replace=givens) _op_ins = clean_inputs + nw_inner op_ins, op_outs = scan_utils.reconstruct_graph(_op_ins, _op_outs) # Reconstruct node nwScan = scan_op.Scan(op_ins, op_outs, op.info) nw_node = nwScan.make_node(*(node.inputs + nw_outer)) env.replace_all_validate(zip(node.outputs, nw_node.outputs), reason="scan_push_computation_out") return True elif to_keep == []: # Nothing in the inner graph should be kept replace_with = {} for idx, out in enumerate(to_replace): if out in local_env.outputs: x = node.outputs[local_env.outputs.index(out)] y = replace_with_out[idx] shape = [y.shape[idx] for idx in xrange(y.ndim)] replace_with[x] = tensor.alloc(y, node.inputs[0], *shape) # We need to add one extra dimension to the outputs env.replace_all_validate(replace_with.items(), reason="scan_push_computation_out") else: return False
def process_node(self, fgraph, node): # this flag tells if there was any change during the last iterations changed = True clean_inputs, clean_outputs = scan_utils.reconstruct_graph( node.op.inputs, node.op.outputs) local_fgraph = gof.FunctionGraph(clean_inputs, clean_outputs) max_iterations = 2 * len(local_fgraph.toposort()) + 3 counts = 0 to_remove = [] to_replace = [] replace_with_in = [] replace_with_out = [] op = node.op # Construct the list of non_sequences to simplify a few things st = op.n_seqs st += int(numpy.sum([len(x) for x in op.tap_array[:(op.n_mit_mot + op.n_mit_sot)]])) st += op.n_sit_sot st += op.n_shared_outs non_seqs = clean_inputs[st:] st = (op.n_seqs + op.n_mit_mot + op.n_mit_sot + op.n_sit_sot + op.n_nit_sot + op.n_shared_outs + 1) outer_non_seqs = node.inputs[st:] assert len(non_seqs) == len(outer_non_seqs) while changed and counts < max_iterations: counts += 1 changed = False for nd in local_fgraph.toposort(): if (numpy.all([(x in non_seqs) or (x.owner in to_remove) or isinstance(x, tensor.Constant) for x in nd.inputs]) and # we can do this because the assumption is that a # viewOp or deepCopyOp will be just at the end of the # function and not somewhere in the middle .. not isinstance(nd.op, theano.compile.ViewOp) and not isinstance(nd.op, theano.compile.DeepCopyOp) and # and we didn't already looked at this node not nd in to_remove): # We have a candidate node to removable # Step 1. Reconstruct it on outside to_remove.append(nd) outside_ins = [] for x in nd.inputs: if x in non_seqs: outside_ins += [outer_non_seqs[non_seqs.index(x)]] elif x in to_replace: outside_ins += [ replace_with_out[to_replace.index(x)]] elif isinstance(x, theano.Constant): outside_ins += [x.clone()] else: raise Exception( ('Error in the `scan_pushout_non_seq_' 'operations`. The optimization tries ' 'to move some computation fron scan ' 'which is not allowed to move. Report ' 'this on theano-users list'), x) outside_ins = [x.type.filter_variable(y) for x,y in zip(nd.inputs, outside_ins)] nw_outer_node = nd.op.make_node(*outside_ins) # Step 2. Create variables for replacements for idx, y in enumerate(nd.outputs): y_place_holder = scan_utils.safe_new(y, '_replace') to_replace += [y] replace_with_in += [y_place_holder] assert type(y) == type(nw_outer_node.outputs[idx]) replace_with_out += [nw_outer_node.outputs[idx]] changed = True if counts >= max_iterations: raise Exception('Error in the `scan_pushout_non_seq_operations`.' ' The optimization exhausted the maximal number ' 'of iterations allowed!') # We need to check all candidate replacements and choose those that # make sense for us # Step 1. which elements of `to_replace` are used by remaining # components of the inner function clean_to_replace = [] clean_replace_with_in = [] clean_replace_with_out = [] existent_nodes = [nd for nd in local_fgraph.toposort() if nd not in to_remove] to_keep = [] for nd in existent_nodes: to_keep += nd.inputs for idx, out in enumerate(to_replace): if out in to_keep and out.owner not in existent_nodes: clean_to_replace += [out] clean_replace_with_in += [replace_with_in[idx]] clean_replace_with_out += [replace_with_out[idx]] if len(clean_to_replace) > 0: # We can finally put an end to all this madness givens = {} nw_outer = [] nw_inner = [] for to_repl, repl_in, repl_out in zip(clean_to_replace, clean_replace_with_in, clean_replace_with_out): if isinstance(repl_out, theano.Constant): repl_in = repl_out.clone() else: nw_inner += [repl_in] nw_outer += [repl_out] givens[to_repl] = repl_in _op_outs = scan_utils.clone(clean_outputs, replace=givens) _op_ins = clean_inputs + nw_inner op_ins, op_outs = scan_utils.reconstruct_graph(_op_ins, _op_outs) # Reconstruct node nwScan = scan_op.Scan(op_ins, op_outs, op.info) nw_node = nwScan.make_node(* (node.inputs + nw_outer)) fgraph.replace_all_validate_remove( zip(node.outputs, nw_node.outputs), remove=[node], reason='scan_push_computation_out') return True elif to_keep == []: # Nothing in the inner graph should be kept replace_with = {} for idx, out in enumerate(to_replace): if out in local_fgraph.outputs: x = node.outputs[local_fgraph.outputs.index(out)] y = replace_with_out[idx] shape = [y.shape[idx] for idx in xrange(y.ndim)] replace_with[x] = tensor.alloc(y, node.inputs[0], *shape) # We need to add one extra dimension to the outputs if replace_with: fgraph.replace_all_validate_remove( replace_with.items(), remove=[node], reason='scan_push_computation_out') else: return False
def scan_merge_inouts(node): if not isinstance(node.op, scan_op.Scan): return False a = scan_args(node.inputs, node.outputs, node.op.inputs, node.op.outputs, node.op.info) inp_equiv = {} if has_duplicates(a.outer_in_seqs): new_outer_seqs = [] new_inner_seqs = [] for out_seq, in_seq in zip(a.outer_in_seqs, a.inner_in_seqs): if out_seq in new_outer_seqs: i = new_outer_seqs.index(out_seq) inp_equiv[in_seq] = new_inner_seqs[i] else: new_outer_seqs.append(out_seq) new_inner_seqs.append(in_seq) a.outer_in_seqs = new_outer_seqs a.inner_in_seqs = new_inner_seqs if has_duplicates(a.outer_in_non_seqs): new_outer_nseqs = [] new_inner_nseqs = [] for out_nseq, in_nseq in zip(a.outer_in_non_seqs, a.inner_in_non_seqs): if out_nseq in new_outer_nseqs: i = new_outer_nseqs.index(out_nseq) inp_equiv[in_nseq] = new_inner_nseqs[i] else: new_outer_nseqs.append(out_nseq) new_inner_nseqs.append(in_nseq) a.outer_in_non_seqs = new_outer_nseqs a.inner_in_non_seqs = new_inner_nseqs if len(inp_equiv) > 0: # do the replacement now. The rest will be left to ScanSaveMem inner_inputs = a.inner_inputs outer_inputs = a.outer_inputs info = a.info if info['as_while']: a_inner_outs = a.inner_outputs + a.cond else: a_inner_outs = a.inner_outputs inner_outputs = scan_utils.clone(a_inner_outs, replace=inp_equiv) op = scan_op.Scan(inner_inputs, inner_outputs, info) outputs = op(*outer_inputs) if not isinstance(outputs, (list, tuple)): outputs = [outputs] na = scan_args(outer_inputs, outputs, op.inputs, op.outputs, op.info) else: na = a # start again left = [] right = [] if has_duplicates(na.outer_in_shared): _left, _right = make_equiv(na.outer_in_shared, na.inner_in_shared) left += _left right += _right if has_duplicates(na.outer_in_sit_sot): _left, _right = make_equiv(na.outer_in_sit_sot, na.inner_in_sit_sot) left += _left right += _right if has_duplicates(na.outer_in_mit_mot): seen = {} for omm, imm, _sl in zip(na.outer_in_mit_mot, na.inner_in_mit_mot, na.mit_mot_in_slices): sl = tuple(_sl) if (omm, sl) in seen: simm = seen[(omm, sl)] left += imm right += simm else: seen[(omm, sl)] = imm if has_duplicates(na.outer_in_mit_sot): seen = {} for oms, ims, _sl in zip(na.outer_in_mit_sot, na.inner_in_mit_sot, na.mit_sot_in_slices): sl = tuple(_sl) if (oms, sl) in seen: sims = seen[(oms, sl)] left += ims right += sims else: seen[(oms, sl)] = ims def map_out(i, o, seen): for si, so in seen: if equal_computations([i], [si], left, right): return so seen.append((i, o)) return o seen = [] na.outer_out_nit_sot = [map_out(i, o, seen) for i, o in zip(na.inner_out_nit_sot, na.outer_out_nit_sot)] seen = [] na.outer_out_sit_sot = [map_out(i, o, seen) for i, o in zip(na.inner_out_sit_sot, na.outer_out_sit_sot)] seen = [] na.outer_out_mit_sot = [map_out(i, o, seen) for i, o in zip(na.inner_out_mit_sot, na.outer_out_mit_sot)] seen = [] new_outer_out_mit_mot = [] for imm, omm, osl in zip(na.inner_out_mit_mot, na.outer_out_mit_mot, na.mit_mot_out_slices): for simm, somm, sosl in seen: if osl == sosl and equal_computations(imm, simm, left, right): new_outer_out_mit_mot.append(somm) break else: seen.append((imm, omm, osl)) new_outer_out_mit_mot.append(omm) na.outer_out_mit_mot = new_outer_out_mit_mot return na.outer_outputs
def scan(fn, sequences=None, outputs_info=None, non_sequences=None, n_steps=None, truncate_gradient=-1, go_backwards=False, mode=None, name=None, options=None, profile=False): """ This function constructs and applies a Scan op to the provided arguments. :param fn: ``fn`` is a function that describes the operations involved in one step of ``scan``. ``fn`` should construct variables describing the output of one iteration step. It should expect as input theano variables representing all the slices of the input sequences and previous values of the outputs, as well as all other arguments given to scan as ``non_sequences``. The order in which scan passes these variables to ``fn`` is the following : * all time slices of the first sequence * all time slices of the second sequence * ... * all time slices of the last sequence * all past slices of the first output * all past slices of the second otuput * ... * all past slices of the last output * all other arguments (the list given as `non_sequences` to scan) The order of the sequences is the same as the one in the list `sequences` given to scan. The order of the outputs is the same as the order of ``output_info``. For any sequence or output the order of the time slices is the same as the one in which they have been given as taps. For example if one writes the following : .. code-block:: python scan(fn, sequences = [ dict(input= Sequence1, taps = [-3,2,-1]) , Sequence2 , dict(input = Sequence3, taps = 3) ] , outputs_info = [ dict(initial = Output1, taps = [-3,-5]) , dict(initial = Output2, taps = None) , Output3 ] , non_sequences = [ Argument1, Argument 2]) ``fn`` should expect the following arguments in this given order: #. ``Sequence1[t-3]`` #. ``Sequence1[t+2]`` #. ``Sequence1[t-1]`` #. ``Sequence2[t]`` #. ``Sequence3[t+3]`` #. ``Output1[t-3]`` #. ``Output1[t-5]`` #. ``Output3[t-1]`` #. ``Argument1`` #. ``Argument2`` The list of ``non_sequences`` can also contain shared variables used in the function, though ``scan`` is able to figure those out on its own so they can be skipped. For the clarity of the code we recommand though to provide them to scan. To some extend ``scan`` can also figure out other ``non sequences`` (not shared) even if not passed to scan (but used by `fn`). A simple example of this would be : .. code-block:: python import theano.tensor as TT W = TT.matrix() W_2 = W**2 def f(x): return TT.dot(x,W_2) The function is expected to return two things. One is a list of outputs ordered in the same order as ``outputs_info``, with the difference that there should be only one output variable per output initial state (even if no tap value is used). Secondly `fn` should return an update dictionary (that tells how to update any shared variable after each iteration step). The dictionary can optionally be given as a list of tuples. There is no constraint on the order of these two list, ``fn`` can return either ``(outputs_list, update_dictionary)`` or ``(update_dictionary, outputs_list)`` or just one of the two (in case the other is empty). To use ``scan`` as a while loop, the user needs to change the function ``fn`` such that also a stopping condition is returned. To do so, he/she needs to wrap the condition in an ``until`` class. The condition should be returned as a third element, for example: .. code-block:: python ... return [y1_t, y2_t], {x:x+1}, theano.scan_module.until(x < 50) Note that a number of steps (considered in here as the maximum number of steps ) is still required even though a condition is passed (and it is used to allocate memory if needed). = {}): :param sequences: ``sequences`` is the list of Theano variables or dictionaries describing the sequences ``scan`` has to iterate over. If a sequence is given as wrapped in a dictionary, then a set of optional information can be provided about the sequence. The dictionary should have the following keys: * ``input`` (*mandatory*) -- Theano variable representing the sequence. * ``taps`` -- Temporal taps of the sequence required by ``fn``. They are provided as a list of integers, where a value ``k`` impiles that at iteration step ``t`` scan will pass to ``fn`` the slice ``t+k``. Default value is ``[0]`` Any Theano variable in the list ``sequences`` is automatically wrapped into a dictionary where ``taps`` is set to ``[0]`` :param outputs_info: ``outputs_info`` is the list of Theano variables or dictionaries describing the initial state of the outputs computed recurrently. When this initial states are given as dictionary optional information can be provided about the output corresponding to these initial states. The dictionary should have the following keys: * ``initial`` -- Theano variable that represents the initial state of a given output. In case the output is not computed recursively (think of a map) and does not require a initial state this field can be skiped. Given that only the previous time step of the output is used by ``fn`` the initial state should have the same shape as the output. If multiple time taps are used, the initial state should have one extra dimension that should cover all the possible taps. For example if we use ``-5``, ``-2`` and ``-1`` as past taps, at step 0, ``fn`` will require (by an abuse of notation) ``output[-5]``, ``output[-2]`` and ``output[-1]``. This will be given by the initial state, which in this case should have the shape (5,)+output.shape. If this variable containing the initial state is called ``init_y`` then ``init_y[0]`` *corresponds to* ``output[-5]``. ``init_y[1]`` *correponds to* ``output[-4]``, ``init_y[2]`` corresponds to ``output[-3]``, ``init_y[3]`` coresponds to ``output[-2]``, ``init_y[4]`` corresponds to ``output[-1]``. While this order might seem strange, it comes natural from splitting an array at a given point. Assume that we have a array ``x``, and we choose ``k`` to be time step ``0``. Then our initial state would be ``x[:k]``, while the output will be ``x[k:]``. Looking at this split, elements in ``x[:k]`` are ordered exactly like those in ``init_y``. * ``taps`` -- Temporal taps of the output that will be pass to ``fn``. They are provided as a list of *negative* integers, where a value ``k`` implies that at iteration step ``t`` scan will pass to ``fn`` the slice ``t+k``. ``scan`` will follow this logic if partial information is given: * If an output is not wrapped in a dictionary, ``scan`` will wrap it in one assuming that you use only the last step of the output (i.e. it makes your tap value list equal to [-1]). * If you wrap an output in a dictionary and you do not provide any taps but you provide an initial state it will assume that you are using only a tap value of -1. * If you wrap an output in a dictionary but you do not provide any initial state, it assumes that you are not using any form of taps. * If you provide a ``None`` instead of a variable or a empty dictionary ``scan`` assumes that you will not use any taps for this output (like for example in case of a map) If ``outputs_info`` is an empty list or None, ``scan`` assumes that no tap is used for any of the outputs. If information is provided just for a subset of the outputs an exception is raised (because there is no convention on how scan should map the provided information to the outputs of ``fn``) :param non_sequences: ``non_sequences`` is the list of arguments that are passed to ``fn`` at each steps. One can opt to exclude variable used in ``fn`` from this list as long as they are part of the computational graph, though for clarity we encourage not to do so. :param n_steps: ``n_steps`` is the number of steps to iterate given as an int or Theano scalar. If any of the input sequences do not have enough elements, scan will raise an error. If the *value is 0* the outputs will have *0 rows*. If the value is negative, ``scan`` will run backwards in time. If the ``go_backwards`` flag is already set and also ``n_steps`` is negative, ``scan`` will run forward in time. If n stpes is not provided, ``scan`` will figure out the amount of steps it should run given its input sequences. :param truncate_gradient: ``truncate_gradient`` is the number of steps to use in truncated BPTT. If you compute gradients through a scan op, they are computed using backpropagation through time. By providing a different value then -1, you choose to use truncated BPTT instead of classical BPTT, where you go for only ``truncate_gradient`` number of steps back in time. :param go_backwards: ``go_backwards`` is a flag indicating if ``scan`` should go backwards through the sequences. If you think of each sequence as indexed by time, making this flag True would mean that ``scan`` goes back in time, namely that for any sequence it starts from the end and goes towards 0. :param name: When profiling ``scan``, it is crucial to provide a name for any instance of ``scan``. The profiler will produce an overall profile of your code as well as profiles for the computation of one step of each instance of ``scan``. The ``name`` of the instance appears in those profiles and can greatly help to disambiguate information. :param mode: It is recommended to leave this argument to None, especially when profiling ``scan`` (otherwise the results are not going to be accurate). If you prefer the computations of one step of ``scan`` to be done differently then the entire function, you can use this parameter to describe how the computations in this loop are done (see ``theano.function`` for details about possible values and their meaning). :param profile: Flag or string. If true, or different from the empty string, a profile object will be created and attached to the inner graph of scan. In case ``profile`` is True, the profile object will have the name of the scan instance, otherwise it will have the passed string. Profile object collect (and print) information only when running the inner graph with the new cvm linker ( with default modes, other linkers this argument is useless) :rtype: tuple :return: tuple of the form (outputs, updates); ``outputs`` is either a Theano variable or a list of Theano variables representing the outputs of ``scan`` (in the same order as in ``outputs_info``). ``updates`` is a subclass of dictionary specifying the update rules for all shared variables used in scan This dictionary should be passed to ``theano.function`` when you compile your function. The change compared to a normal dictionary is that we validate that keys are SharedVariable and addition of those dictionary are validated to be consistent. """ # Note : see the internal documentation of the scan op for naming # conventions and all other details if options is None: options = {} rvals = scan_utils.canonical_arguments(sequences, outputs_info, non_sequences, go_backwards, n_steps) inputs, states_and_outputs_info, parameters, T = rvals # If we provided a known number of steps ( before compilation) # and if that number is 1 or -1, then we can skip the Scan Op, # and just apply the inner function once # To do that we check here to see the nature of n_steps T_value = None if isinstance(n_steps, (float, int)): T_value = int(n_steps) else: try: T_value = opt.get_constant_value(n_steps) except (TypeError, AttributeError): T_value = None if T_value in (1, -1): return one_step_scan(fn, inputs, states_and_outputs_info, parameters, truncate_gradient) # 1. Variable representing the current time step t = scalar_shared(numpy.int64(0), name='t') # 2. Allocate memory for the states of scan. mintaps = [] lengths = [] for pos, arg_info in enumerate(states_and_outputs_info): if arg_info.get('taps', None) == [-1]: mintaps.append(1) lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) arg_info['initial'] = scan_utils.expand(tensor.unbroadcast( tensor.shape_padleft(arg_info['initial']), 0), T) elif arg_info.get('taps', None): if numpy.any(numpy.array(arg_info.get('taps', [])) > 0): # Make sure we do not have requests for future values of a # sequence we can not provide such values raise ValueError('Can not use future taps of outputs', arg_info) mintap = abs(numpy.min(arg_info['taps'])) lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) mintaps.append(mintap) arg_info['initial'] = scan_utils.expand( arg_info['initial'][:mintap], T) else: mintaps.append(0) lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) # 3. Generate arguments for the function passed to scan. This will # function will return the outputs that need to be computed at every # timesteps inputs_slices = [input[t] for input in inputs] states_slices = [] for n, state in enumerate(states_and_outputs_info): # Check if it is actually a state and not an output if mintaps[n] != 0: for k in state['taps']: states_slices.append( state['initial'][(t + mintaps[n] + k) % lengths[n]]) # 4. Construct outputs that are to be computed by the inner # function of scan args = inputs_slices + states_slices + parameters cond, states_and_outputs, updates = \ scan_utils.get_updates_and_outputs(fn(*args)) # User is allowed to provide no information if it only behaves like a # map if (len(states_and_outputs) != len(states_and_outputs_info) and len(states_and_outputs_info) == 0): mintaps = [0] * len(states_and_outputs) # 5. Construct the scan op # 5.1 Construct list of shared variables with updates (those that # can be treated as states (i.e. of TensorType) and those that can not # (like Random States) if cond is not None: _cond = [cond] else: _cond = [] rvals = rebuild_collect_shared( states_and_outputs + _cond, updates=updates, rebuild_strict=True, copy_inputs_over=True, no_default_updates=False) # extracting the arguments input_variables, cloned_outputs, other_rval = rvals clone_d, update_d, update_expr, shared_inputs = other_rval additional_input_states = [] additional_output_states = [] additional_lengths = [] additional_mintaps = [] original_numeric_shared_variables = [] non_numeric_input_states = [] non_numeric_output_states = [] original_non_numeric_shared_variables = [] pos = len(lengths) for sv in shared_inputs: if sv in update_d: if isinstance(sv, (TensorVariable, TensorSharedVariable)): # We can treat it as a sit sot nw_state = scan_utils.expand( tensor.unbroadcast(tensor.shape_padleft(sv), 0), T) additional_lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) pos = pos + 1 additional_mintaps.append(1) additional_input_states.append(nw_state) additional_output_states.append( scan_utils.clone(tensor.set_subtensor( nw_state[(t + 1) % additional_lengths[-1]], update_d[sv]))) original_numeric_shared_variables.append(sv) else: non_numeric_input_states.append(sv) non_numeric_output_states.append(update_d[sv]) original_non_numeric_shared_variables.append(sv) # Replace shared variables in the update _additional_output_states = [] replace = {} for sv, buf in zip(original_numeric_shared_variables, additional_input_states): replace[sv] = buf[t] for out in additional_output_states: _additional_output_states.append( scan_utils.clone(out, replace=replace)) additional_output_states = _additional_output_states # 5.2 Collect inputs/outputs of the inner function inputs = [] outputs = [] for n, mintap in enumerate(mintaps): if mintap != 0: input_state = states_and_outputs_info[n]['initial'] inputs.append(input_state) outputs.append( tensor.set_subtensor( input_state[(t + mintap) % lengths[n]], states_and_outputs[n])) else: mem_buffer = scan_utils.allocate_memory( T, states_and_outputs_info[n], states_and_outputs[n]) inputs.append(output) outputs.append( tensor.set_subtensor(output[t % lengths[n]], states_and_outputs[n])) inputs.extend(additional_input_states) outputs.extend(additional_output_states) lengths.extend(additional_lengths) mintaps.extend(additional_mintaps) inputs.extend(non_numeric_input_states) outputs.extend(non_numeric_output_states) all_other_inputs = gof.graph.inputs(outputs) parameters = [x for x in all_other_inputs if (x not in inputs and x not in lengths and x is not t and isinstance(x, gof.Variable) and not isinstance(x, gof.Constant))] inputs.extend(parameters) # 5.3 Construct the the options dictionary options['name'] = name options['profile'] = profile options['mode'] = mode options['inplace'] = False options['gpu'] = False options['truncate_gradient'] = truncate_gradient options['hash_inner_graph'] = 0 # 5.4 Construct the ScanOp instance local_op = scan_op.ScanOp(inputs=inputs, outputs=outputs, lengths=lengths, switches=[], mintaps=mintaps, index=t, options=options, as_repeatUntil=cond) # Note that we get here all the outputs followed by the update rules to # the shared variables we had in our scan # we know that we have (in this given order): # * len(states_and_outputs) real outputs # * len(additional_input_states) updates for numeric shared variable # * len(non_numeric_input_states) updates for non numeric shared # variables scan_inputs = [T] + inputs scan_outputs_update_rules = scan_utils.to_list(local_op(*scan_inputs)) # 5.5 Collect outputs and add permutation object scan_outputs = [] for pos in xrange(len(states_and_outputs)): out = scan_utils.ScanPermutation(mintaps[pos])( scan_outputs_update_rules[pos], t) scan_outputs.append(out[mintaps[pos]:]) # 5.6 Construct updates dictionary update_rules = scan_outputs_update_rules[len(states_and_outputs):] updates = {} for v, u in izip(original_numeric_shared_variables, update_rules[:len(additional_input_states)]): updates[v] = u[-1] for v, u in izip(original_non_numeric_shared_variables, update_rules[len(additional_input_states):]): updates[v] = u # Step 5.7 We are done and can return everything back to the user return scan_outputs, updates
def scan(fn, sequences=None, outputs_info=None, non_sequences=None, n_steps=None, truncate_gradient=-1, go_backwards=False, mode=None, name=None, options=None, profile=False): """ This function constructs and applies a Scan op to the provided arguments. :param fn: ``fn`` is a function that describes the operations involved in one step of ``scan``. ``fn`` should construct variables describing the output of one iteration step. It should expect as input theano variables representing all the slices of the input sequences and previous values of the outputs, as well as all other arguments given to scan as ``non_sequences``. The order in which scan passes these variables to ``fn`` is the following : * all time slices of the first sequence * all time slices of the second sequence * ... * all time slices of the last sequence * all past slices of the first output * all past slices of the second otuput * ... * all past slices of the last output * all other arguments (the list given as `non_sequences` to scan) The order of the sequences is the same as the one in the list `sequences` given to scan. The order of the outputs is the same as the order of ``outputs_info``. For any sequence or output the order of the time slices is the same as the one in which they have been given as taps. For example if one writes the following : .. code-block:: python scan(fn, sequences = [ dict(input= Sequence1, taps = [-3,2,-1]) , Sequence2 , dict(input = Sequence3, taps = 3) ] , outputs_info = [ dict(initial = Output1, taps = [-3,-5]) , dict(initial = Output2, taps = None) , Output3 ] , non_sequences = [ Argument1, Argument 2]) ``fn`` should expect the following arguments in this given order: #. ``Sequence1[t-3]`` #. ``Sequence1[t+2]`` #. ``Sequence1[t-1]`` #. ``Sequence2[t]`` #. ``Sequence3[t+3]`` #. ``Output1[t-3]`` #. ``Output1[t-5]`` #. ``Output3[t-1]`` #. ``Argument1`` #. ``Argument2`` The list of ``non_sequences`` can also contain shared variables used in the function, though ``scan`` is able to figure those out on its own so they can be skipped. For the clarity of the code we recommend though to provide them to scan. To some extend ``scan`` can also figure out other ``non sequences`` (not shared) even if not passed to scan (but used by `fn`). A simple example of this would be : .. code-block:: python import theano.tensor as TT W = TT.matrix() W_2 = W**2 def f(x): return TT.dot(x,W_2) The function is expected to return two things. One is a list of outputs ordered in the same order as ``outputs_info``, with the difference that there should be only one output variable per output initial state (even if no tap value is used). Secondly `fn` should return an update dictionary (that tells how to update any shared variable after each iteration step). The dictionary can optionally be given as a list of tuples. There is no constraint on the order of these two list, ``fn`` can return either ``(outputs_list, update_dictionary)`` or ``(update_dictionary, outputs_list)`` or just one of the two (in case the other is empty). To use ``scan`` as a while loop, the user needs to change the function ``fn`` such that also a stopping condition is returned. To do so, he/she needs to wrap the condition in an ``until`` class. The condition should be returned as a third element, for example: .. code-block:: python ... return [y1_t, y2_t], {x:x+1}, theano.scan_module.until(x < 50) Note that a number of steps (considered in here as the maximum number of steps ) is still required even though a condition is passed (and it is used to allocate memory if needed). = {}): :param sequences: ``sequences`` is the list of Theano variables or dictionaries describing the sequences ``scan`` has to iterate over. If a sequence is given as wrapped in a dictionary, then a set of optional information can be provided about the sequence. The dictionary should have the following keys: * ``input`` (*mandatory*) -- Theano variable representing the sequence. * ``taps`` -- Temporal taps of the sequence required by ``fn``. They are provided as a list of integers, where a value ``k`` impiles that at iteration step ``t`` scan will pass to ``fn`` the slice ``t+k``. Default value is ``[0]`` Any Theano variable in the list ``sequences`` is automatically wrapped into a dictionary where ``taps`` is set to ``[0]`` :param outputs_info: ``outputs_info`` is the list of Theano variables or dictionaries describing the initial state of the outputs computed recurrently. When this initial states are given as dictionary optional information can be provided about the output corresponding to these initial states. The dictionary should have the following keys: * ``initial`` -- Theano variable that represents the initial state of a given output. In case the output is not computed recursively (think of a map) and does not require a initial state this field can be skiped. Given that only the previous time step of the output is used by ``fn`` the initial state should have the same shape as the output. If multiple time taps are used, the initial state should have one extra dimension that should cover all the possible taps. For example if we use ``-5``, ``-2`` and ``-1`` as past taps, at step 0, ``fn`` will require (by an abuse of notation) ``output[-5]``, ``output[-2]`` and ``output[-1]``. This will be given by the initial state, which in this case should have the shape (5,)+output.shape. If this variable containing the initial state is called ``init_y`` then ``init_y[0]`` *corresponds to* ``output[-5]``. ``init_y[1]`` *correponds to* ``output[-4]``, ``init_y[2]`` corresponds to ``output[-3]``, ``init_y[3]`` coresponds to ``output[-2]``, ``init_y[4]`` corresponds to ``output[-1]``. While this order might seem strange, it comes natural from splitting an array at a given point. Assume that we have a array ``x``, and we choose ``k`` to be time step ``0``. Then our initial state would be ``x[:k]``, while the output will be ``x[k:]``. Looking at this split, elements in ``x[:k]`` are ordered exactly like those in ``init_y``. * ``taps`` -- Temporal taps of the output that will be pass to ``fn``. They are provided as a list of *negative* integers, where a value ``k`` implies that at iteration step ``t`` scan will pass to ``fn`` the slice ``t+k``. ``scan`` will follow this logic if partial information is given: * If an output is not wrapped in a dictionary, ``scan`` will wrap it in one assuming that you use only the last step of the output (i.e. it makes your tap value list equal to [-1]). * If you wrap an output in a dictionary and you do not provide any taps but you provide an initial state it will assume that you are using only a tap value of -1. * If you wrap an output in a dictionary but you do not provide any initial state, it assumes that you are not using any form of taps. * If you provide a ``None`` instead of a variable or a empty dictionary ``scan`` assumes that you will not use any taps for this output (like for example in case of a map) If ``outputs_info`` is an empty list or None, ``scan`` assumes that no tap is used for any of the outputs. If information is provided just for a subset of the outputs an exception is raised (because there is no convention on how scan should map the provided information to the outputs of ``fn``) :param non_sequences: ``non_sequences`` is the list of arguments that are passed to ``fn`` at each steps. One can opt to exclude variable used in ``fn`` from this list as long as they are part of the computational graph, though for clarity we encourage not to do so. :param n_steps: ``n_steps`` is the number of steps to iterate given as an int or Theano scalar. If any of the input sequences do not have enough elements, scan will raise an error. If the *value is 0* the outputs will have *0 rows*. If the value is negative, ``scan`` will run backwards in time. If the ``go_backwards`` flag is already set and also ``n_steps`` is negative, ``scan`` will run forward in time. If n stpes is not provided, ``scan`` will figure out the amount of steps it should run given its input sequences. :param truncate_gradient: ``truncate_gradient`` is the number of steps to use in truncated BPTT. If you compute gradients through a scan op, they are computed using backpropagation through time. By providing a different value then -1, you choose to use truncated BPTT instead of classical BPTT, where you go for only ``truncate_gradient`` number of steps back in time. :param go_backwards: ``go_backwards`` is a flag indicating if ``scan`` should go backwards through the sequences. If you think of each sequence as indexed by time, making this flag True would mean that ``scan`` goes back in time, namely that for any sequence it starts from the end and goes towards 0. :param name: When profiling ``scan``, it is crucial to provide a name for any instance of ``scan``. The profiler will produce an overall profile of your code as well as profiles for the computation of one step of each instance of ``scan``. The ``name`` of the instance appears in those profiles and can greatly help to disambiguate information. :param mode: It is recommended to leave this argument to None, especially when profiling ``scan`` (otherwise the results are not going to be accurate). If you prefer the computations of one step of ``scan`` to be done differently then the entire function, you can use this parameter to describe how the computations in this loop are done (see ``theano.function`` for details about possible values and their meaning). :param profile: Flag or string. If true, or different from the empty string, a profile object will be created and attached to the inner graph of scan. In case ``profile`` is True, the profile object will have the name of the scan instance, otherwise it will have the passed string. Profile object collect (and print) information only when running the inner graph with the new cvm linker ( with default modes, other linkers this argument is useless) :rtype: tuple :return: tuple of the form (outputs, updates); ``outputs`` is either a Theano variable or a list of Theano variables representing the outputs of ``scan`` (in the same order as in ``outputs_info``). ``updates`` is a subclass of dictionary specifying the update rules for all shared variables used in scan This dictionary should be passed to ``theano.function`` when you compile your function. The change compared to a normal dictionary is that we validate that keys are SharedVariable and addition of those dictionary are validated to be consistent. """ # Note : see the internal documentation of the scan op for naming # conventions and all other details if options is None: options = {} rvals = scan_utils.canonical_arguments(sequences, outputs_info, non_sequences, go_backwards, n_steps) inputs, states_and_outputs_info, parameters, T = rvals # If we provided a known number of steps ( before compilation) # and if that number is 1 or -1, then we can skip the Scan Op, # and just apply the inner function once # To do that we check here to see the nature of n_steps T_value = None if isinstance(n_steps, (float, int)): T_value = int(n_steps) else: try: T_value = opt.get_scalar_constant_value(n_steps) except (TypeError, AttributeError): T_value = None if T_value in (1, -1): return one_step_scan(fn, inputs, states_and_outputs_info, parameters, truncate_gradient) # 1. Variable representing the current time step t = scalar_shared(numpy.int64(0), name='t') # 2. Allocate memory for the states of scan. mintaps = [] lengths = [] for pos, arg_info in enumerate(states_and_outputs_info): if arg_info.get('taps', None) == [-1]: mintaps.append(1) lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) arg_info['initial'] = scan_utils.expand(tensor.unbroadcast( tensor.shape_padleft(arg_info['initial']), 0), T) elif arg_info.get('taps', None): if numpy.any(numpy.array(arg_info.get('taps', [])) > 0): # Make sure we do not have requests for future values of a # sequence we can not provide such values raise ValueError('Can not use future taps of outputs', arg_info) mintap = abs(numpy.min(arg_info['taps'])) lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) mintaps.append(mintap) arg_info['initial'] = scan_utils.expand( arg_info['initial'][:mintap], T) else: mintaps.append(0) lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) # 3. Generate arguments for the function passed to scan. This will # function will return the outputs that need to be computed at every # timesteps inputs_slices = [input[t] for input in inputs] states_slices = [] for n, state in enumerate(states_and_outputs_info): # Check if it is actually a state and not an output if mintaps[n] != 0: for k in state['taps']: states_slices.append( state['initial'][(t + mintaps[n] + k) % lengths[n]]) # 4. Construct outputs that are to be computed by the inner # function of scan args = inputs_slices + states_slices + parameters cond, states_and_outputs, updates = \ scan_utils.get_updates_and_outputs(fn(*args)) # User is allowed to provide no information if it only behaves like a # map if (len(states_and_outputs) != len(states_and_outputs_info) and len(states_and_outputs_info) == 0): mintaps = [0] * len(states_and_outputs) # 5. Construct the scan op # 5.1 Construct list of shared variables with updates (those that # can be treated as states (i.e. of TensorType) and those that can not # (like Random States) if cond is not None: _cond = [cond] else: _cond = [] rvals = rebuild_collect_shared( states_and_outputs + _cond, updates=updates, rebuild_strict=True, copy_inputs_over=True, no_default_updates=False) # extracting the arguments input_variables, cloned_outputs, other_rval = rvals clone_d, update_d, update_expr, shared_inputs = other_rval additional_input_states = [] additional_output_states = [] additional_lengths = [] additional_mintaps = [] original_numeric_shared_variables = [] non_numeric_input_states = [] non_numeric_output_states = [] original_non_numeric_shared_variables = [] pos = len(lengths) for sv in shared_inputs: if sv in update_d: if isinstance(sv, (TensorVariable, TensorSharedVariable)): # We can treat it as a sit sot nw_state = scan_utils.expand( tensor.unbroadcast(tensor.shape_padleft(sv), 0), T) additional_lengths.append(scalar_shared(numpy.int64(0), name='l%d' % pos)) pos = pos + 1 additional_mintaps.append(1) additional_input_states.append(nw_state) additional_output_states.append( scan_utils.clone(tensor.set_subtensor( nw_state[(t + 1) % additional_lengths[-1]], update_d[sv]))) original_numeric_shared_variables.append(sv) else: non_numeric_input_states.append(sv) non_numeric_output_states.append(update_d[sv]) original_non_numeric_shared_variables.append(sv) # Replace shared variables in the update _additional_output_states = [] replace = {} for sv, buf in zip(original_numeric_shared_variables, additional_input_states): replace[sv] = buf[t] for out in additional_output_states: _additional_output_states.append( scan_utils.clone(out, replace=replace)) additional_output_states = _additional_output_states # 5.2 Collect inputs/outputs of the inner function inputs = [] outputs = [] for n, mintap in enumerate(mintaps): if mintap != 0: input_state = states_and_outputs_info[n]['initial'] inputs.append(input_state) outputs.append( tensor.set_subtensor( input_state[(t + mintap) % lengths[n]], states_and_outputs[n])) else: mem_buffer = scan_utils.allocate_memory( T, states_and_outputs_info[n], states_and_outputs[n]) inputs.append(output) outputs.append( tensor.set_subtensor(output[t % lengths[n]], states_and_outputs[n])) inputs.extend(additional_input_states) outputs.extend(additional_output_states) lengths.extend(additional_lengths) mintaps.extend(additional_mintaps) inputs.extend(non_numeric_input_states) outputs.extend(non_numeric_output_states) all_other_inputs = gof.graph.inputs(outputs) parameters = [x for x in all_other_inputs if (x not in inputs and x not in lengths and x is not t and isinstance(x, gof.Variable) and not isinstance(x, gof.Constant))] inputs.extend(parameters) # 5.3 Construct the the options dictionary options['name'] = name options['profile'] = profile options['mode'] = mode options['inplace'] = False options['gpu'] = False options['truncate_gradient'] = truncate_gradient options['hash_inner_graph'] = 0 # 5.4 Construct the ScanOp instance local_op = scan_op.ScanOp(inputs=inputs, outputs=outputs, lengths=lengths, switches=[], mintaps=mintaps, index=t, options=options, as_repeatUntil=cond) # Note that we get here all the outputs followed by the update rules to # the shared variables we had in our scan # we know that we have (in this given order): # * len(states_and_outputs) real outputs # * len(additional_input_states) updates for numeric shared variable # * len(non_numeric_input_states) updates for non numeric shared # variables scan_inputs = [T] + inputs scan_outputs_update_rules = scan_utils.to_list(local_op(*scan_inputs)) # 5.5 Collect outputs and add permutation object scan_outputs = [] for pos in xrange(len(states_and_outputs)): out = scan_utils.ScanPermutation(mintaps[pos])( scan_outputs_update_rules[pos], t) scan_outputs.append(out[mintaps[pos]:]) # 5.6 Construct updates dictionary update_rules = scan_outputs_update_rules[len(states_and_outputs):] updates = {} for v, u in izip(original_numeric_shared_variables, update_rules[:len(additional_input_states)]): updates[v] = u[-1] for v, u in izip(original_non_numeric_shared_variables, update_rules[len(additional_input_states):]): updates[v] = u # Step 5.7 We are done and can return everything back to the user return scan_outputs, updates
## ### Step 4. Compile the dummy function ## # We can now compile a dummy function just to see what shared variable # we have and what are their update rules (note that the user has # the option not to pass the shared variable to scan, so we need to # pick them manually and add them to scan) # make the compilation as fast as possible by not applying any # optimization or conversion to C [ note this region is not important # for performance so we can do stuff as unoptimal as we wish ] # extract still missing inputs (there still might be so) and add them # as non sequences at the end of our args fake_nonseqs = [x.type() for x in non_seqs] fake_outputs = scan_utils.clone(outputs, replace=dict(zip(non_seqs, fake_nonseqs))) all_inputs = itertools.ifilter( lambda x: (isinstance(x, gof.Variable) and not isinstance( x, SharedVariable) and not isinstance(x, gof.Constant)), gof.graph.inputs(fake_outputs)) extra_inputs = filter(lambda x: x not in args + fake_nonseqs, all_inputs) non_seqs += extra_inputs ## Note we do not use all_inputs directly since the order of variables ## in args is quite important dummy_args += extra_inputs dummy_outs = outputs if condition is not None: dummy_outs.append(condition) dummy_f = function(dummy_args, dummy_outs,