<s> import theano <EOL> import copy <EOL> from theano import Op <EOL> from theano . gof import local_optimizer <EOL> from theano . sandbox . cuda import cuda_available , GpuOp <EOL> from theano . sandbox . cuda . basic_ops import gpu_flatten <EOL> from theano . tensor . extra_ops import CumsumOp <EOL> if cuda_available : <EOL> from theano . sandbox . cuda import CudaNdarrayType <EOL> from theano . sandbox . cuda . basic_ops import host_from_gpu , gpu_from_host , HostFromGpu <EOL> from theano . sandbox . cuda import register_opt as register_gpu_opt <EOL> class GpuCumsum ( CumsumOp , GpuOp ) : <EOL> """<STR_LIT>""" <EOL> SUPPORTED_NDIMS = <NUM_LIT:3> <EOL> __props__ = ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def __init__ ( self , axis ) : <EOL> self . axis = axis <EOL> self . max_threads_dim0 = None <EOL> self . max_grid_size1 = None <EOL> self . max_grid_size2 = None <EOL> perform = Op . perform <EOL> def make_node ( self , x ) : <EOL> assert x . dtype == '<STR_LIT>' <EOL> if not isinstance ( x . type , CudaNdarrayType ) : <EOL> raise TypeError ( '<STR_LIT>' , x ) <EOL> if x . ndim > GpuCumsum . SUPPORTED_NDIMS : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> if self . axis >= x . ndim or self . axis < - x . ndim : <EOL> raise ValueError ( '<STR_LIT>' . format ( self . axis ) ) <EOL> return theano . Apply ( self , [ x ] , [ x . type ( ) ] ) <EOL> def make_thunk ( self , node , storage_map , compute_map , no_recycling ) : <EOL> node_ = copy . copy ( node ) <EOL> assert node . op is node_ . op <EOL> if node_ . op . max_threads_dim0 is None or node_ . op . max_grid_size1 is None or node_ . op . max_grid_size2 is None : <EOL> cuda = theano . sandbox . cuda <EOL> device_id = cuda . use . device_number <EOL> if device_id is None : <EOL> cuda . use ( "<STR_LIT>" , <EOL> force = False , <EOL> default_to_move_computation_to_gpu = False , <EOL> move_shared_float32_to_gpu = False , <EOL> enable_cuda = False , <EOL> test_driver = True ) <EOL> device_id = cuda . use . device_number <EOL> cuda_ndarray = theano . sandbox . cuda . cuda_ndarray . cuda_ndarray <EOL> prop = cuda_ndarray . device_properties ( device_id ) <EOL> node_ . op . max_threads_dim0 = prop [ '<STR_LIT>' ] <EOL> node_ . op . max_grid_size1 = prop [ '<STR_LIT>' ] <EOL> node_ . op . max_grid_size2 = prop [ '<STR_LIT>' ] <EOL> return super ( GpuCumsum , node_ . op ) . make_thunk ( node_ , storage_map , <EOL> compute_map , no_recycling ) <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , self . axis ) <EOL> def c_code_cache_version ( self ) : <EOL> return ( <NUM_LIT:9> , ) <EOL> def c_support_code_apply ( self , node , nodename ) : <EOL> return """<STR_LIT>""" % locals ( ) <EOL> def c_code ( self , node , nodename , inames , onames , sub ) : <EOL> x , = inames <EOL> z , = onames <EOL> axis = self . axis if self . axis is not None else <NUM_LIT:0> <EOL> fail = sub [ '<STR_LIT>' ] <EOL> max_threads_dim0 = self . max_threads_dim0 <EOL> max_grid_size1 = self . max_grid_size1 <EOL> max_grid_size2 = self . max_grid_size2 <EOL> if max_threads_dim0 is None or max_grid_size1 is None or max_grid_size2 is None : <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> code = """<STR_LIT>""" % locals ( ) <EOL> return code <EOL> def values_eq_approx_high_tol ( a , b ) : <EOL> """<STR_LIT>""" <EOL> rtol = None <EOL> if a . size > <NUM_LIT> : <EOL> rtol = <NUM_LIT> <EOL> return CudaNdarrayType . values_eq_approx ( a , b , rtol = rtol ) <EOL> @ register_gpu_opt ( ) <EOL> @ local_optimizer ( [ CumsumOp ] ) <EOL> def use_gpu_cumsum ( node ) : <EOL> if type ( node . op ) is CumsumOp and node . inputs [ <NUM_LIT:0> ] . dtype == '<STR_LIT>' and node . inputs [ <NUM_LIT:0> ] . owner and isinstance ( node . inputs [ <NUM_LIT:0> ] . owner . op , HostFromGpu ) : <EOL> axis = node . op . axis <EOL> x = node . inputs [ <NUM_LIT:0> ] <EOL> if axis is not None and x . ndim > GpuCumsum . SUPPORTED_NDIMS : <EOL> return None <EOL> x = gpu_from_host ( x ) <EOL> if axis is None and x . ndim > <NUM_LIT:1> : <EOL> x = gpu_flatten ( x ) <EOL> if axis is None : <EOL> axis = <NUM_LIT:0> <EOL> ret = host_from_gpu ( GpuCumsum ( axis ) ( x ) ) <EOL> ret . tag . values_eq_approx = values_eq_approx_high_tol <EOL> return [ ret ] </s>
