<s> from __future__ import absolute_import , print_function , division <EOL> from collections import defaultdict <EOL> from six import iteritems <EOL> from theano . gof . graph import list_of_nodes <EOL> from theano . compat import cmp <EOL> def memodict ( f ) : <EOL> """<STR_LIT>""" <EOL> class memodict ( defaultdict ) : <EOL> def __missing__ ( self , key ) : <EOL> ret = self [ key ] = f ( key ) <EOL> return ret <EOL> return memodict ( ) . __getitem__ <EOL> def make_depends ( ) : <EOL> @ memodict <EOL> def depends ( pair ) : <EOL> """<STR_LIT>""" <EOL> a , b = pair <EOL> return ( any ( bout in a . inputs for bout in b . outputs ) or <EOL> any ( depends ( ( ainp . owner , b ) ) for ainp in a . inputs <EOL> if ainp . owner ) ) <EOL> return depends <EOL> def make_dependence_cmp ( ) : <EOL> """<STR_LIT>""" <EOL> depends = make_depends ( ) <EOL> def dependence ( a , b ) : <EOL> """<STR_LIT>""" <EOL> if depends ( ( a , b ) ) : <EOL> return <NUM_LIT:1> <EOL> if depends ( ( b , a ) ) : <EOL> return - <NUM_LIT:1> <EOL> return <NUM_LIT:0> <EOL> return dependence <EOL> def reverse_dict ( d ) : <EOL> """<STR_LIT>""" <EOL> result = { } <EOL> for key in d : <EOL> for val in d [ key ] : <EOL> result [ val ] = result . get ( val , tuple ( ) ) + ( key , ) <EOL> return result <EOL> def _toposort ( edges ) : <EOL> """<STR_LIT>""" <EOL> incoming_edges = reverse_dict ( edges ) <EOL> incoming_edges = dict ( ( k , set ( val ) ) <EOL> for k , val in iteritems ( incoming_edges ) ) <EOL> S = set ( ( v for v in edges if v not in incoming_edges ) ) <EOL> L = [ ] <EOL> while S : <EOL> n = S . pop ( ) <EOL> L . append ( n ) <EOL> for m in edges . get ( n , ( ) ) : <EOL> assert n in incoming_edges [ m ] <EOL> incoming_edges [ m ] . remove ( n ) <EOL> if not incoming_edges [ m ] : <EOL> S . add ( m ) <EOL> if any ( incoming_edges . get ( v , None ) for v in edges ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return L <EOL> def posort ( l , * cmps ) : <EOL> """<STR_LIT>""" <EOL> comes_before = dict ( ( a , set ( ) ) for a in l ) <EOL> comes_after = dict ( ( a , set ( ) ) for a in l ) <EOL> def add_links ( a , b ) : <EOL> comes_after [ a ] . add ( b ) <EOL> comes_after [ a ] . update ( comes_after [ b ] ) <EOL> for c in comes_before [ a ] : <EOL> comes_after [ c ] . update ( comes_after [ a ] ) <EOL> comes_before [ b ] . add ( a ) <EOL> comes_before [ b ] . update ( comes_before [ a ] ) <EOL> for c in comes_after [ b ] : <EOL> comes_before [ c ] . update ( comes_before [ b ] ) <EOL> def check ( ) : <EOL> """<STR_LIT>""" <EOL> for a in l : <EOL> for b in l : <EOL> assert not ( b in comes_after [ a ] and a in comes_after [ b ] ) <EOL> for cmp_fn in cmps : <EOL> for a in l : <EOL> for b in l : <EOL> if cmp_fn ( a , b ) < <NUM_LIT:0> : <EOL> if b not in comes_before [ a ] and b not in comes_after [ a ] : <EOL> add_links ( a , b ) <EOL> return _toposort ( comes_after ) <EOL> def sort_apply_nodes ( inputs , outputs , cmps ) : <EOL> """<STR_LIT>""" <EOL> return posort ( list_of_nodes ( inputs , outputs ) , * cmps ) <EOL> def sort_schedule_fn ( * cmps ) : <EOL> """<STR_LIT>""" <EOL> dependence = make_dependence_cmp ( ) <EOL> cmps = ( dependence , ) + cmps <EOL> def schedule ( fgraph ) : <EOL> """<STR_LIT>""" <EOL> return sort_apply_nodes ( fgraph . inputs , fgraph . outputs , cmps ) <EOL> return schedule <EOL> def key_to_cmp ( key ) : <EOL> def key_cmp ( a , b ) : <EOL> return cmp ( key ( a ) , key ( b ) ) <EOL> return key_cmp </s>
