<s> import sys <EOL> import numpy as np <EOL> from pyll import Apply <EOL> from pyll import as_apply <EOL> from pyll import dfs <EOL> from pyll import toposort <EOL> from pyll import scope <EOL> from pyll import stochastic <EOL> from pyll import clone_merge <EOL> stoch = stochastic . implicit_stochastic_symbols <EOL> def ERR ( msg ) : <EOL> print >> sys . stderr , '<STR_LIT>' , msg <EOL> @ scope . define_pure <EOL> def vchoice_split ( idxs , choices , n_options ) : <EOL> rval = [ [ ] for ii in range ( n_options ) ] <EOL> if len ( idxs ) != len ( choices ) : <EOL> raise ValueError ( '<STR_LIT>' , <EOL> ( len ( idxs ) , len ( choices ) ) ) <EOL> for ii , cc in zip ( idxs , choices ) : <EOL> rval [ cc ] . append ( ii ) <EOL> return rval <EOL> @ scope . define_pure <EOL> def vchoice_merge ( idxs , choices , * vals ) : <EOL> rval = [ ] <EOL> assert len ( idxs ) == len ( choices ) <EOL> for idx , ch in zip ( idxs , choices ) : <EOL> vi , vv = vals [ ch ] <EOL> rval . append ( vv [ list ( vi ) . index ( idx ) ] ) <EOL> return rval <EOL> @ scope . define_pure <EOL> def idxs_map ( idxs , cmd , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if <NUM_LIT:0> : <EOL> for ii , ( idxs_ii , vals_ii ) in enumerate ( args ) : <EOL> for jj in idxs : assert jj in idxs_ii <EOL> for kw , ( idxs_kw , vals_kw ) in kwargs . items ( ) : <EOL> for jj in idxs : assert jj in idxs_kw <EOL> args_imap = [ ] <EOL> for idxs_j , vals_j in args : <EOL> if len ( idxs_j ) : <EOL> args_imap . append ( dict ( zip ( idxs_j , vals_j ) ) ) <EOL> else : <EOL> args_imap . append ( { } ) <EOL> kwargs_imap = { } <EOL> for kw , ( idxs_j , vals_j ) in kwargs . items ( ) : <EOL> if len ( idxs_j ) : <EOL> kwargs_imap [ kw ] = dict ( zip ( idxs_j , vals_j ) ) <EOL> else : <EOL> kwargs_imap [ kw ] = { } <EOL> f = scope . _impls [ cmd ] <EOL> rval = [ ] <EOL> for ii in idxs : <EOL> try : <EOL> args_nn = [ arg_imap [ ii ] for arg_imap in args_imap ] <EOL> except : <EOL> ERR ( '<STR_LIT>' % cmd ) <EOL> ERR ( '<STR_LIT>' % ii ) <EOL> ERR ( '<STR_LIT>' % str ( arg_imap ) ) <EOL> ERR ( '<STR_LIT>' % str ( args_imap ) ) <EOL> raise <EOL> try : <EOL> kwargs_nn = dict ( [ ( kw , arg_imap [ ii ] ) <EOL> for kw , arg_imap in kwargs_imap . items ( ) ] ) <EOL> except : <EOL> ERR ( '<STR_LIT>' % cmd ) <EOL> ERR ( '<STR_LIT>' % ii ) <EOL> ERR ( '<STR_LIT>' % kw ) <EOL> ERR ( '<STR_LIT>' % str ( arg_imap ) ) <EOL> raise <EOL> try : <EOL> rval_nn = f ( * args_nn , ** kwargs_nn ) <EOL> except : <EOL> ERR ( '<STR_LIT>' % cmd ) <EOL> raise <EOL> rval . append ( rval_nn ) <EOL> return rval <EOL> @ scope . define_pure <EOL> def idxs_take ( idxs , vals , which ) : <EOL> """<STR_LIT>""" <EOL> assert len ( idxs ) == len ( vals ) <EOL> table = dict ( zip ( idxs , vals ) ) <EOL> return np . asarray ( [ table [ w ] for w in which ] ) <EOL> @ scope . define_pure <EOL> def uniq ( lst ) : <EOL> s = set ( ) <EOL> rval = [ ] <EOL> for l in lst : <EOL> if id ( l ) not in s : <EOL> s . add ( id ( l ) ) <EOL> rval . append ( l ) <EOL> return rval <EOL> def vectorize_stochastic ( orig ) : <EOL> if orig . name == '<STR_LIT>' and orig . pos_args [ <NUM_LIT:1> ] . _obj in stoch : <EOL> idxs = orig . pos_args [ <NUM_LIT:0> ] <EOL> dist = orig . pos_args [ <NUM_LIT:1> ] . _obj <EOL> def foo ( arg ) : <EOL> assert arg . name == '<STR_LIT>' <EOL> assert len ( arg . pos_args ) == <NUM_LIT:2> <EOL> arg_vals = arg . pos_args [ <NUM_LIT:1> ] <EOL> if arg_vals . name == '<STR_LIT>' : <EOL> if arg_vals . arg [ '<STR_LIT>' ] . name == '<STR_LIT>' : <EOL> if arg_vals . arg [ '<STR_LIT>' ] . inputs ( ) [ <NUM_LIT:0> ] . name == '<STR_LIT>' : <EOL> repeated_thing = arg_vals . arg [ '<STR_LIT>' ] . inputs ( ) [ <NUM_LIT:0> ] . inputs ( ) [ <NUM_LIT:1> ] <EOL> return repeated_thing <EOL> if arg . pos_args [ <NUM_LIT:0> ] is idxs : <EOL> return arg_vals <EOL> else : <EOL> raise NotImplementedError ( ) <EOL> new_pos_args = [ foo ( arg ) for arg in orig . pos_args [ <NUM_LIT:2> : ] ] <EOL> new_named_args = [ [ aname , foo ( arg ) ] <EOL> for aname , arg in orig . named_args ] <EOL> vnode = Apply ( dist , new_pos_args , new_named_args , o_len = None ) <EOL> n_times = scope . len ( idxs ) <EOL> if '<STR_LIT:size>' in dict ( vnode . named_args ) : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> vnode . named_args . append ( [ '<STR_LIT:size>' , n_times ] ) <EOL> return vnode <EOL> else : <EOL> return orig <EOL> def replace_repeat_stochastic ( expr , return_memo = False ) : <EOL> nodes = dfs ( expr ) <EOL> memo = { } <EOL> for ii , orig in enumerate ( nodes ) : <EOL> if orig . name == '<STR_LIT>' and orig . pos_args [ <NUM_LIT:1> ] . _obj in stoch : <EOL> idxs = orig . pos_args [ <NUM_LIT:0> ] <EOL> dist = orig . pos_args [ <NUM_LIT:1> ] . _obj <EOL> def foo ( arg ) : <EOL> assert arg . name == '<STR_LIT>' <EOL> assert len ( arg . pos_args ) == <NUM_LIT:2> <EOL> arg_vals = arg . pos_args [ <NUM_LIT:1> ] <EOL> if ( arg_vals . name == '<STR_LIT>' <EOL> and arg_vals . inputs ( ) [ <NUM_LIT:0> ] . name == '<STR_LIT>' ) : <EOL> repeated_thing = arg_vals . inputs ( ) [ <NUM_LIT:0> ] . inputs ( ) [ <NUM_LIT:1> ] <EOL> return repeated_thing <EOL> else : <EOL> if arg . pos_args [ <NUM_LIT:0> ] is idxs : <EOL> return arg_vals <EOL> else : <EOL> raise NotImplementedError ( ) <EOL> new_pos_args = [ foo ( arg ) for arg in orig . pos_args [ <NUM_LIT:2> : ] ] <EOL> new_named_args = [ [ aname , foo ( arg ) ] <EOL> for aname , arg in orig . named_args ] <EOL> vnode = Apply ( dist , new_pos_args , new_named_args , None ) <EOL> n_times = scope . len ( idxs ) <EOL> if '<STR_LIT:size>' in dict ( vnode . named_args ) : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> vnode . named_args . append ( [ '<STR_LIT:size>' , n_times ] ) <EOL> for client in nodes [ ii + <NUM_LIT:1> : ] : <EOL> client . replace_input ( orig , vnode ) <EOL> if expr is orig : <EOL> expr = vnode <EOL> memo [ orig ] = vnode <EOL> if return_memo : <EOL> return expr , memo <EOL> else : <EOL> return expr <EOL> class VectorizeHelper ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , expr , expr_idxs , build = True ) : <EOL> self . expr = expr <EOL> self . expr_idxs = expr_idxs <EOL> self . dfs_nodes = dfs ( expr ) <EOL> self . params = { } <EOL> for ii , node in enumerate ( self . dfs_nodes ) : <EOL> if node . name == '<STR_LIT>' : <EOL> label = node . arg [ '<STR_LIT:label>' ] . obj <EOL> self . params [ label ] = node . arg [ '<STR_LIT>' ] <EOL> self . idxs_memo = { } <EOL> self . take_memo = { } <EOL> self . v_expr = self . build_idxs_vals ( expr , expr_idxs ) <EOL> self . assert_integrity_idxs_take ( ) <EOL> def assert_integrity_idxs_take ( self ) : <EOL> idxs_memo = self . idxs_memo <EOL> take_memo = self . take_memo <EOL> after = dfs ( self . expr ) <EOL> assert after == self . dfs_nodes <EOL> assert set ( idxs_memo . keys ( ) ) == set ( take_memo . keys ( ) ) <EOL> for node in idxs_memo : <EOL> idxs = idxs_memo [ node ] <EOL> assert idxs . name == '<STR_LIT>' <EOL> vals = take_memo [ node ] [ <NUM_LIT:0> ] . pos_args [ <NUM_LIT:1> ] <EOL> for take in take_memo [ node ] : <EOL> assert take . name == '<STR_LIT>' <EOL> assert [ idxs , vals ] == take . pos_args [ : <NUM_LIT:2> ] <EOL> def build_idxs_vals ( self , node , wanted_idxs ) : <EOL> """<STR_LIT>""" <EOL> checkpoint_asserts = False <EOL> def checkpoint ( ) : <EOL> if checkpoint_asserts : <EOL> self . assert_integrity_idxs_take ( ) <EOL> if node in self . idxs_memo : <EOL> toposort ( self . idxs_memo [ node ] ) <EOL> if node in self . take_memo : <EOL> for take in self . take_memo [ node ] : <EOL> toposort ( take ) <EOL> checkpoint ( ) <EOL> assert wanted_idxs != self . idxs_memo . get ( node ) <EOL> if node . name == '<STR_LIT>' : <EOL> return self . build_idxs_vals ( node . arg [ '<STR_LIT>' ] , wanted_idxs ) <EOL> elif node . name == '<STR_LIT>' : <EOL> return self . build_idxs_vals ( node . arg [ '<STR_LIT>' ] , wanted_idxs ) <EOL> elif node . name == '<STR_LIT>' : <EOL> if node in self . idxs_memo : <EOL> all_idxs , all_vals = self . take_memo [ node ] [ <NUM_LIT:0> ] . pos_args [ : <NUM_LIT:2> ] <EOL> wanted_vals = scope . idxs_take ( all_idxs , all_vals , wanted_idxs ) <EOL> self . take_memo [ node ] . append ( wanted_vals ) <EOL> checkpoint ( ) <EOL> else : <EOL> all_idxs = self . expr_idxs <EOL> n_times = scope . len ( all_idxs ) <EOL> all_idxs = scope . array_union ( all_idxs ) <EOL> self . idxs_memo [ node ] = all_idxs <EOL> all_vals = scope . asarray ( scope . repeat ( n_times , node ) ) <EOL> wanted_vals = scope . idxs_take ( all_idxs , all_vals , wanted_idxs ) <EOL> assert node not in self . take_memo <EOL> self . take_memo [ node ] = [ wanted_vals ] <EOL> checkpoint ( ) <EOL> return wanted_vals <EOL> elif node . name == '<STR_LIT>' : <EOL> if ( node in self . idxs_memo <EOL> and wanted_idxs in self . idxs_memo [ node ] . pos_args ) : <EOL> all_idxs , all_vals = self . take_memo [ node ] [ <NUM_LIT:0> ] . pos_args [ : <NUM_LIT:2> ] <EOL> wanted_vals = scope . idxs_take ( all_idxs , all_vals , wanted_idxs ) <EOL> self . take_memo [ node ] . append ( wanted_vals ) <EOL> checkpoint ( ) <EOL> else : <EOL> if node in self . idxs_memo : <EOL> all_idxs = self . idxs_memo [ node ] <EOL> assert all_idxs . name == '<STR_LIT>' <EOL> all_idxs . pos_args . append ( wanted_idxs ) <EOL> else : <EOL> all_idxs = scope . array_union ( wanted_idxs ) <EOL> choice = node . pos_args [ <NUM_LIT:0> ] <EOL> all_choices = self . build_idxs_vals ( choice , all_idxs ) <EOL> options = node . pos_args [ <NUM_LIT:1> : ] <EOL> args_idxs = scope . vchoice_split ( all_idxs , all_choices , <EOL> len ( options ) ) <EOL> all_vals = scope . vchoice_merge ( all_idxs , all_choices ) <EOL> for opt_ii , idxs_ii in zip ( options , args_idxs ) : <EOL> all_vals . pos_args . append ( <EOL> as_apply ( [ <EOL> idxs_ii , <EOL> self . build_idxs_vals ( opt_ii , idxs_ii ) , <EOL> ] ) ) <EOL> wanted_vals = scope . idxs_take ( <EOL> all_idxs , <EOL> all_vals , <EOL> wanted_idxs ) <EOL> if node in self . idxs_memo : <EOL> assert self . idxs_memo [ node ] . name == '<STR_LIT>' <EOL> self . idxs_memo [ node ] . pos_args . append ( wanted_idxs ) <EOL> for take in self . take_memo [ node ] : <EOL> assert take . name == '<STR_LIT>' <EOL> take . pos_args [ <NUM_LIT:1> ] = all_vals <EOL> self . take_memo [ node ] . append ( wanted_vals ) <EOL> else : <EOL> self . idxs_memo [ node ] = all_idxs <EOL> self . take_memo [ node ] = [ wanted_vals ] <EOL> checkpoint ( ) <EOL> else : <EOL> if ( node in self . idxs_memo <EOL> and wanted_idxs in self . idxs_memo [ node ] . pos_args ) : <EOL> for take in self . take_memo [ node ] : <EOL> if take . pos_args [ <NUM_LIT:2> ] == wanted_idxs : <EOL> return take <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> else : <EOL> checkpoint ( ) <EOL> if node in self . idxs_memo : <EOL> all_idxs = self . idxs_memo [ node ] <EOL> else : <EOL> all_idxs = scope . array_union ( wanted_idxs ) <EOL> checkpoint ( ) <EOL> all_vals = scope . idxs_map ( all_idxs , node . name ) <EOL> for ii , aa in enumerate ( node . pos_args ) : <EOL> all_vals . pos_args . append ( as_apply ( [ <EOL> all_idxs , self . build_idxs_vals ( aa , all_idxs ) ] ) ) <EOL> checkpoint ( ) <EOL> for ii , ( nn , aa ) in enumerate ( node . named_args ) : <EOL> all_vals . named_args . append ( [ nn , as_apply ( [ <EOL> all_idxs , self . build_idxs_vals ( aa , all_idxs ) ] ) ] ) <EOL> checkpoint ( ) <EOL> all_vals = vectorize_stochastic ( all_vals ) <EOL> checkpoint ( ) <EOL> wanted_vals = scope . idxs_take ( <EOL> all_idxs , <EOL> all_vals , <EOL> wanted_idxs ) <EOL> if node in self . idxs_memo : <EOL> assert self . idxs_memo [ node ] . name == '<STR_LIT>' <EOL> self . idxs_memo [ node ] . pos_args . append ( wanted_idxs ) <EOL> toposort ( self . idxs_memo [ node ] ) <EOL> for take in self . take_memo [ node ] : <EOL> assert take . name == '<STR_LIT>' <EOL> take . pos_args [ <NUM_LIT:1> ] = all_vals <EOL> self . take_memo [ node ] . append ( wanted_vals ) <EOL> else : <EOL> self . idxs_memo [ node ] = all_idxs <EOL> self . take_memo [ node ] = [ wanted_vals ] <EOL> checkpoint ( ) <EOL> return wanted_vals <EOL> def idxs_by_label ( self ) : <EOL> return dict ( [ ( name , self . idxs_memo [ node ] ) <EOL> for name , node in self . params . items ( ) ] ) <EOL> def vals_by_label ( self ) : <EOL> return dict ( [ ( name , self . take_memo [ node ] [ <NUM_LIT:0> ] . pos_args [ <NUM_LIT:1> ] ) <EOL> for name , node in self . params . items ( ) ] ) </s>
