<s> import numpy <EOL> from theano import gof <EOL> from theano . gof import Constant , Generic , Op <EOL> from theano . gof . sched import key_to_cmp <EOL> from theano . tensor import tensor <EOL> import theano <EOL> class LoadFromDisk ( Op ) : <EOL> """<STR_LIT>""" <EOL> __props__ = ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) <EOL> def __init__ ( self , dtype , broadcastable , mmap_mode = None ) : <EOL> self . dtype = numpy . dtype ( dtype ) <EOL> self . broadcastable = broadcastable <EOL> if mmap_mode not in ( None , '<STR_LIT:c>' ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % mmap_mode ) <EOL> self . mmap_mode = mmap_mode <EOL> def make_node ( self , path ) : <EOL> if isinstance ( path , str ) : <EOL> path = Constant ( Generic ( ) , path ) <EOL> return gof . Apply ( self , [ path ] , [ tensor ( self . dtype , <EOL> broadcastable = self . broadcastable ) ] ) <EOL> def perform ( self , node , inp , out ) : <EOL> path = inp [ <NUM_LIT:0> ] <EOL> if ( path . split ( '<STR_LIT:.>' ) [ - <NUM_LIT:1> ] == '<STR_LIT>' ) : <EOL> raise ValueError ( "<STR_LIT>" % path ) <EOL> result = numpy . load ( path , mmap_mode = self . mmap_mode ) <EOL> if result . dtype != self . dtype : <EOL> raise TypeError ( "<STR_LIT>" % <EOL> ( self . dtype , result . dtype ) ) <EOL> out [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] = result <EOL> def __str__ ( self ) : <EOL> return ( "<STR_LIT>" % <EOL> ( self . dtype , self . broadcastable , self . mmap_mode ) ) <EOL> def load ( path , dtype , broadcastable , mmap_mode = None ) : <EOL> """<STR_LIT>""" <EOL> return LoadFromDisk ( dtype , broadcastable , mmap_mode ) ( path ) <EOL> try : <EOL> from mpi4py import MPI <EOL> except ImportError : <EOL> mpi_enabled = False <EOL> else : <EOL> comm = MPI . COMM_WORLD <EOL> mpi_enabled = True <EOL> class MPIRecv ( Op ) : <EOL> """<STR_LIT>""" <EOL> __props__ = ( "<STR_LIT:source>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) <EOL> def __init__ ( self , source , tag , shape , dtype ) : <EOL> self . source = source <EOL> self . tag = tag <EOL> self . shape = shape <EOL> self . dtype = numpy . dtype ( dtype ) <EOL> self . broadcastable = ( False , ) * len ( shape ) <EOL> def make_node ( self ) : <EOL> return gof . Apply ( self , [ ] , [ theano . Variable ( Generic ( ) ) , <EOL> tensor ( self . dtype , <EOL> broadcastable = self . broadcastable ) ] ) <EOL> def perform ( self , node , inp , out ) : <EOL> data = numpy . zeros ( self . shape , dtype = self . dtype ) <EOL> request = comm . Irecv ( data , self . source , self . tag ) <EOL> out [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] = request <EOL> out [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] = data <EOL> def __str__ ( self ) : <EOL> return ( "<STR_LIT>" % <EOL> ( self . source , self . tag , self . shape , self . dtype ) ) <EOL> def infer_shape ( self , node , shapes ) : <EOL> return [ None , self . shape ] <EOL> def do_constant_folding ( self , node ) : <EOL> return False <EOL> class MPIRecvWait ( Op ) : <EOL> """<STR_LIT>""" <EOL> __props__ = ( "<STR_LIT>" , ) <EOL> def __init__ ( self , tag ) : <EOL> self . tag = tag <EOL> def make_node ( self , request , data ) : <EOL> return gof . Apply ( self , [ request , data ] , <EOL> [ tensor ( data . dtype , <EOL> broadcastable = data . broadcastable ) ] ) <EOL> def perform ( self , node , inp , out ) : <EOL> request = inp [ <NUM_LIT:0> ] <EOL> data = inp [ <NUM_LIT:1> ] <EOL> request . wait ( ) <EOL> out [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] = data <EOL> def infer_shape ( self , node , shapes ) : <EOL> return [ shapes [ <NUM_LIT:1> ] ] <EOL> view_map = { <NUM_LIT:0> : [ <NUM_LIT:1> ] } <EOL> class MPISend ( Op ) : <EOL> """<STR_LIT>""" <EOL> __props__ = ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> def __init__ ( self , dest , tag ) : <EOL> self . dest = dest <EOL> self . tag = tag <EOL> def make_node ( self , data ) : <EOL> return gof . Apply ( self , [ data ] , <EOL> [ theano . Variable ( Generic ( ) ) , data . type ( ) ] ) <EOL> view_map = { <NUM_LIT:1> : [ <NUM_LIT:0> ] } <EOL> def perform ( self , node , inp , out ) : <EOL> data = inp [ <NUM_LIT:0> ] <EOL> request = comm . Isend ( data , self . dest , self . tag ) <EOL> out [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] = request <EOL> out [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] = data <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" % ( self . dest , self . tag ) <EOL> class MPISendWait ( Op ) : <EOL> """<STR_LIT>""" <EOL> __props__ = ( "<STR_LIT>" , ) <EOL> def __init__ ( self , tag ) : <EOL> self . tag = tag <EOL> def make_node ( self , request , data ) : <EOL> return gof . Apply ( self , [ request , data ] , <EOL> [ theano . Variable ( Generic ( ) ) ] ) <EOL> def perform ( self , node , inp , out ) : <EOL> request = inp [ <NUM_LIT:0> ] <EOL> request . wait ( ) <EOL> out [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] = True <EOL> def isend ( var , dest , tag ) : <EOL> """<STR_LIT>""" <EOL> return MPISend ( dest , tag ) ( var ) <EOL> def send ( var , dest , tag ) : <EOL> """<STR_LIT>""" <EOL> return MPISendWait ( tag ) ( * isend ( var , dest , tag ) ) <EOL> def irecv ( shape , dtype , source , tag ) : <EOL> """<STR_LIT>""" <EOL> return MPIRecv ( source , tag , shape , dtype ) ( ) <EOL> def recv ( shape , dtype , source , tag ) : <EOL> """<STR_LIT>""" <EOL> return MPIRecvWait ( tag ) ( * irecv ( shape , dtype , source , tag ) ) <EOL> def mpi_send_wait_key ( a ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( a . op , ( MPIRecvWait , MPISendWait ) ) : <EOL> return <NUM_LIT:1> <EOL> if isinstance ( a . op , ( MPIRecv , MPISend ) ) : <EOL> return - <NUM_LIT:1> <EOL> return <NUM_LIT:0> <EOL> def mpi_tag_key ( a ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( a . op , ( MPISend , MPIRecv , MPIRecvWait , MPISendWait ) ) : <EOL> return a . op . tag <EOL> else : <EOL> return <NUM_LIT:0> <EOL> mpi_send_wait_cmp = key_to_cmp ( mpi_send_wait_key ) <EOL> mpi_tag_cmp = key_to_cmp ( mpi_tag_key ) <EOL> mpi_keys = ( mpi_send_wait_key , mpi_tag_key ) <EOL> mpi_cmps = ( mpi_send_wait_cmp , mpi_tag_cmp ) </s>
