<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import , division <EOL> import atexit <EOL> import numbers <EOL> from string import Template <EOL> import pycuda . driver as drv <EOL> import pycuda . gpuarray as gpuarray <EOL> import pycuda . elementwise as elementwise <EOL> import pycuda . reduction as reduction <EOL> import pycuda . scan as scan <EOL> import pycuda . tools as tools <EOL> from pycuda . tools import context_dependent_memoize , dtype_to_ctype <EOL> from pycuda . compiler import SourceModule <EOL> from pytools import memoize <EOL> import numpy as np <EOL> from . import cuda <EOL> from . import cublas <EOL> try : <EOL> from . import cula <EOL> _has_cula = True <EOL> except ( ImportError , OSError ) : <EOL> _has_cula = False <EOL> try : <EOL> from . import cusolver <EOL> _has_cusolver = True <EOL> except ( ImportError , OSError ) : <EOL> _has_cusolver = False <EOL> try : <EOL> from . import magma <EOL> _has_magma = True <EOL> except ( ImportError , OSError ) : <EOL> _has_magma = False <EOL> isdoubletype = lambda x : True if x == np . float64 or x == np . complex128 else False <EOL> isdoubletype . __doc__ = """<STR_LIT>""" <EOL> iscomplextype = lambda x : True if x == np . complex64 or x == np . complex128 else False <EOL> iscomplextype . __doc__ = """<STR_LIT>""" <EOL> def init_device ( n = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> drv . init ( ) <EOL> dev = drv . Device ( n ) <EOL> return dev <EOL> def init_context ( dev ) : <EOL> """<STR_LIT>""" <EOL> ctx = dev . make_context ( ) <EOL> atexit . register ( ctx . pop ) <EOL> return ctx <EOL> def done_context ( ctx ) : <EOL> """<STR_LIT>""" <EOL> for i in xrange ( len ( atexit . _exithandlers ) ) : <EOL> if atexit . _exithandlers [ i ] [ <NUM_LIT:0> ] == ctx . pop : <EOL> del atexit . _exithandlers [ i ] <EOL> break <EOL> ctx . detach ( ) <EOL> global _global_cublas_handle <EOL> _global_cublas_handle = None <EOL> global _global_cusolver_handle <EOL> _global_cusolver_handle = None <EOL> global _global_cublas_allocator <EOL> _global_cublas_allocator = None <EOL> def init ( allocator = drv . mem_alloc ) : <EOL> """<STR_LIT>""" <EOL> global _global_cublas_handle , _global_cublas_allocator <EOL> if not _global_cublas_handle : <EOL> from . import cublas <EOL> _global_cublas_handle = cublas . cublasCreate ( ) <EOL> if _global_cublas_allocator is None : <EOL> _global_cublas_allocator = allocator <EOL> global _global_cusolver_handle <EOL> if not _global_cusolver_handle : <EOL> from . import cusolver <EOL> _global_cusolver_handle = cusolver . cusolverDnCreate ( ) <EOL> if _has_cula : <EOL> cula . culaInitialize ( ) <EOL> if _has_magma : <EOL> magma . magma_init ( ) <EOL> def shutdown ( ) : <EOL> """<STR_LIT>""" <EOL> global _global_cublas_handle <EOL> if _global_cublas_handle : <EOL> from . import cublas <EOL> cublas . cublasDestroy ( _global_cublas_handle ) <EOL> _global_cublas_handle = None <EOL> global _global_cusolver_handle <EOL> if _global_cusolver_handle : <EOL> from . import cusolver <EOL> cusolver . cusolverDnDestroy ( _global_cusolver_handle ) <EOL> _global_cusolver_handle = None <EOL> if _has_cula : <EOL> cula . culaShutdown ( ) <EOL> def get_compute_capability ( dev ) : <EOL> """<STR_LIT>""" <EOL> return np . float ( '<STR_LIT:.>' . join ( [ str ( i ) for i in <EOL> dev . compute_capability ( ) ] ) ) <EOL> def get_current_device ( ) : <EOL> """<STR_LIT>""" <EOL> return drv . Device ( cuda . cudaGetDevice ( ) ) <EOL> @ memoize <EOL> def get_dev_attrs ( dev ) : <EOL> """<STR_LIT>""" <EOL> attrs = dev . get_attributes ( ) <EOL> return [ attrs [ drv . device_attribute . MAX_THREADS_PER_BLOCK ] , <EOL> ( attrs [ drv . device_attribute . MAX_BLOCK_DIM_X ] , <EOL> attrs [ drv . device_attribute . MAX_BLOCK_DIM_Y ] , <EOL> attrs [ drv . device_attribute . MAX_BLOCK_DIM_Z ] ) , <EOL> ( attrs [ drv . device_attribute . MAX_GRID_DIM_X ] , <EOL> attrs [ drv . device_attribute . MAX_GRID_DIM_Y ] , <EOL> attrs [ drv . device_attribute . MAX_GRID_DIM_Z ] ) ] <EOL> iceil = lambda n : int ( np . ceil ( n ) ) <EOL> @ memoize <EOL> def select_block_grid_sizes ( dev , data_shape , threads_per_block = None ) : <EOL> """<STR_LIT>""" <EOL> if np . isscalar ( data_shape ) : <EOL> data_shape = ( data_shape , ) <EOL> N = int ( np . prod ( data_shape ) ) <EOL> max_threads_per_block , max_block_dim , max_grid_dim = get_dev_attrs ( dev ) <EOL> if threads_per_block is not None : <EOL> if threads_per_block > max_threads_per_block : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> else : <EOL> max_threads_per_block = threads_per_block <EOL> blocks_needed = iceil ( N / float ( max_threads_per_block ) ) <EOL> if blocks_needed <= max_grid_dim [ <NUM_LIT:0> ] : <EOL> return ( max_threads_per_block , <NUM_LIT:1> , <NUM_LIT:1> ) , ( blocks_needed , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> elif blocks_needed > max_grid_dim [ <NUM_LIT:0> ] and blocks_needed <= max_grid_dim [ <NUM_LIT:0> ] * max_grid_dim [ <NUM_LIT:1> ] : <EOL> return ( max_threads_per_block , <NUM_LIT:1> , <NUM_LIT:1> ) , ( max_grid_dim [ <NUM_LIT:0> ] , iceil ( blocks_needed / float ( max_grid_dim [ <NUM_LIT:0> ] ) ) , <NUM_LIT:1> ) <EOL> elif blocks_needed > max_grid_dim [ <NUM_LIT:0> ] * max_grid_dim [ <NUM_LIT:1> ] and blocks_needed <= max_grid_dim [ <NUM_LIT:0> ] * max_grid_dim [ <NUM_LIT:1> ] * max_grid_dim [ <NUM_LIT:2> ] : <EOL> return ( max_threads_per_block , <NUM_LIT:1> , <NUM_LIT:1> ) , ( max_grid_dim [ <NUM_LIT:0> ] , max_grid_dim [ <NUM_LIT:1> ] , <EOL> iceil ( blocks_needed / float ( max_grid_dim [ <NUM_LIT:0> ] * max_grid_dim [ <NUM_LIT:1> ] ) ) ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> def zeros ( shape , dtype , order = '<STR_LIT:C>' , allocator = drv . mem_alloc ) : <EOL> """<STR_LIT>""" <EOL> out = gpuarray . GPUArray ( shape , dtype , allocator , order = order ) <EOL> out . fill ( <NUM_LIT:0> ) <EOL> return out <EOL> def zeros_like ( a ) : <EOL> """<STR_LIT>""" <EOL> out = gpuarray . GPUArray ( a . shape , a . dtype , drv . mem_alloc , <EOL> strides = a . strides ) <EOL> out . fill ( <NUM_LIT:0> ) <EOL> return out <EOL> def ones ( shape , dtype , order = '<STR_LIT:C>' , allocator = drv . mem_alloc ) : <EOL> """<STR_LIT>""" <EOL> out = gpuarray . GPUArray ( shape , dtype , allocator , order = order ) <EOL> out . fill ( <NUM_LIT:1> ) <EOL> return out <EOL> def ones_like ( a ) : <EOL> """<STR_LIT>""" <EOL> out = gpuarray . GPUArray ( a . shape , a . dtype , <EOL> a . allocator , strides = a . strides ) <EOL> out . fill ( <NUM_LIT:1> ) <EOL> return out <EOL> def inf ( shape , dtype , order = '<STR_LIT:C>' , allocator = drv . mem_alloc ) : <EOL> """<STR_LIT>""" <EOL> out = gpuarray . GPUArray ( shape , dtype , allocator , order = order ) <EOL> out . fill ( np . inf ) <EOL> return out <EOL> def maxabs ( x_gpu ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> func = maxabs . cache [ x_gpu . dtype ] <EOL> except KeyError : <EOL> ctype = tools . dtype_to_ctype ( x_gpu . dtype ) <EOL> use_double = int ( x_gpu . dtype in [ np . float64 , np . complex128 ] ) <EOL> ret_type = np . float64 if use_double else np . float32 <EOL> func = reduction . ReductionKernel ( ret_type , neutral = "<STR_LIT:0>" , <EOL> reduce_expr = "<STR_LIT>" , <EOL> map_expr = "<STR_LIT>" , <EOL> arguments = "<STR_LIT>" . format ( ctype = ctype ) ) <EOL> maxabs . cache [ x_gpu . dtype ] = func <EOL> return func ( x_gpu ) <EOL> maxabs . cache = { } <EOL> def cumsum ( x_gpu ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> func = cumsum . cache [ x_gpu . dtype ] <EOL> except KeyError : <EOL> func = scan . InclusiveScanKernel ( x_gpu . dtype , '<STR_LIT>' , <EOL> preamble = '<STR_LIT>' ) <EOL> cumsum . cache [ x_gpu . dtype ] = func <EOL> return func ( x_gpu ) <EOL> cumsum . cache = { } <EOL> def diff ( x_gpu ) : <EOL> """<STR_LIT>""" <EOL> y_gpu = gpuarray . empty ( len ( x_gpu ) - <NUM_LIT:1> , x_gpu . dtype ) <EOL> try : <EOL> func = diff . cache [ x_gpu . dtype ] <EOL> except KeyError : <EOL> ctype = tools . dtype_to_ctype ( x_gpu . dtype ) <EOL> func = elementwise . ElementwiseKernel ( "<STR_LIT>" . format ( ctype = ctype ) , <EOL> "<STR_LIT>" ) <EOL> diff . cache [ x_gpu . dtype ] = func <EOL> func ( x_gpu , y_gpu ) <EOL> return y_gpu <EOL> diff . cache = { } <EOL> num_types = [ np . typeDict [ t ] for t in np . typecodes [ '<STR_LIT>' ] + np . typecodes [ '<STR_LIT>' ] ] <EOL> num_nbytes = dict ( ( np . dtype ( t ) , t ( <NUM_LIT:1> ) . nbytes ) for t in num_types ) <EOL> def set_realloc ( x_gpu , data ) : <EOL> """<STR_LIT>""" <EOL> if x_gpu . shape != data . shape or x_gpu . size != data . size or x_gpu . strides != data . strides or x_gpu . dtype != data . dtype : <EOL> x_gpu . gpudata . free ( ) <EOL> nbytes = num_nbytes [ data . dtype ] <EOL> x_gpu . gpudata = drv . mem_alloc ( nbytes * data . size ) <EOL> x_gpu . shape = data . shape <EOL> x_gpu . size = data . size <EOL> x_gpu . strides = data . strides <EOL> x_gpu . dtype = data . dtype <EOL> x_gpu . set ( data ) <EOL> def get_by_index ( src_gpu , ind ) : <EOL> """<STR_LIT>""" <EOL> assert len ( np . shape ( ind ) ) == <NUM_LIT:1> <EOL> assert issubclass ( ind . dtype . type , numbers . Integral ) <EOL> N = len ( ind ) <EOL> if not isinstance ( ind , gpuarray . GPUArray ) : <EOL> ind = gpuarray . to_gpu ( ind ) <EOL> dest_gpu = gpuarray . empty ( N , dtype = src_gpu . dtype ) <EOL> if N == <NUM_LIT:0> : <EOL> return dest_gpu <EOL> try : <EOL> func = get_by_index . cache [ ( src_gpu . dtype , ind . dtype ) ] <EOL> except KeyError : <EOL> data_ctype = tools . dtype_to_ctype ( src_gpu . dtype ) <EOL> ind_ctype = tools . dtype_to_ctype ( ind . dtype ) <EOL> v = "<STR_LIT>" . format ( data_ctype = data_ctype , ind_ctype = ind_ctype ) <EOL> func = elementwise . ElementwiseKernel ( v , "<STR_LIT>" ) <EOL> get_by_index . cache [ ( src_gpu . dtype , ind . dtype ) ] = func <EOL> func ( dest_gpu , ind , src_gpu , range = slice ( <NUM_LIT:0> , N , <NUM_LIT:1> ) ) <EOL> return dest_gpu <EOL> get_by_index . cache = { } <EOL> def set_by_index ( dest_gpu , ind , src_gpu , ind_which = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> assert len ( np . shape ( ind ) ) == <NUM_LIT:1> <EOL> assert dest_gpu . dtype == src_gpu . dtype <EOL> assert issubclass ( ind . dtype . type , numbers . Integral ) <EOL> N = len ( ind ) <EOL> if N == <NUM_LIT:0> : <EOL> return <EOL> if ind_which == '<STR_LIT>' : <EOL> assert N == len ( src_gpu ) <EOL> elif ind_which == '<STR_LIT:src>' : <EOL> assert N == len ( dest_gpu ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if not isinstance ( ind , gpuarray . GPUArray ) : <EOL> ind = gpuarray . to_gpu ( ind ) <EOL> try : <EOL> func = set_by_index . cache [ ( dest_gpu . dtype , ind . dtype , ind_which ) ] <EOL> except KeyError : <EOL> data_ctype = tools . dtype_to_ctype ( dest_gpu . dtype ) <EOL> ind_ctype = tools . dtype_to_ctype ( ind . dtype ) <EOL> v = "<STR_LIT>" . format ( data_ctype = data_ctype , ind_ctype = ind_ctype ) <EOL> if ind_which == '<STR_LIT>' : <EOL> func = elementwise . ElementwiseKernel ( v , "<STR_LIT>" ) <EOL> else : <EOL> func = elementwise . ElementwiseKernel ( v , "<STR_LIT>" ) <EOL> set_by_index . cache [ ( dest_gpu . dtype , ind . dtype , ind_which ) ] = func <EOL> func ( dest_gpu , ind , src_gpu , range = slice ( <NUM_LIT:0> , N , <NUM_LIT:1> ) ) <EOL> set_by_index . cache = { } <EOL> @ context_dependent_memoize <EOL> def _get_binaryop_vecmat_kernel ( dtype , binary_op ) : <EOL> template = Template ( """<STR_LIT>""" ) <EOL> cache_dir = None <EOL> ctype = dtype_to_ctype ( dtype ) <EOL> tmpl = template . substitute ( type = ctype , binary_op = binary_op ) <EOL> mod = SourceModule ( tmpl ) <EOL> add_row_vec_kernel = mod . get_function ( '<STR_LIT>' ) <EOL> add_col_vec_kernel = mod . get_function ( '<STR_LIT>' ) <EOL> return add_row_vec_kernel , add_col_vec_kernel <EOL> def binaryop_matvec ( binary_op , x_gpu , a_gpu , axis = None , out = None , stream = None ) : <EOL> """<STR_LIT>""" <EOL> if axis is None : <EOL> if len ( a_gpu . shape ) == <NUM_LIT:1> : <EOL> if a_gpu . shape [ <NUM_LIT:0> ] == x_gpu . shape [ <NUM_LIT:1> ] : <EOL> axis = <NUM_LIT:1> <EOL> else : <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % ( x_gpu . shape , a_gpu . shape ) ) <EOL> elif a_gpu . shape [ <NUM_LIT:1> ] == x_gpu . shape [ <NUM_LIT:1> ] : <EOL> axis = <NUM_LIT:1> <EOL> elif a_gpu . shape [ <NUM_LIT:0> ] == x_gpu . shape [ <NUM_LIT:0> ] : <EOL> axis = <NUM_LIT:0> <EOL> else : <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % ( x_gpu . shape , a_gpu . shape ) ) <EOL> else : <EOL> if axis < <NUM_LIT:0> : <EOL> axis += <NUM_LIT:2> <EOL> if axis > <NUM_LIT:1> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if binary_op not in [ '<STR_LIT:+>' , '<STR_LIT:->' , '<STR_LIT:/>' , '<STR_LIT:*>' , '<STR_LIT:%>' ] : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> row_kernel , col_kernel = _get_binaryop_vecmat_kernel ( x_gpu . dtype , binary_op ) <EOL> n , m = np . int32 ( x_gpu . shape [ <NUM_LIT:0> ] ) , np . int32 ( x_gpu . shape [ <NUM_LIT:1> ] ) <EOL> block = ( <NUM_LIT> , <NUM_LIT> , <NUM_LIT:1> ) <EOL> gridx = int ( n // block [ <NUM_LIT:0> ] + <NUM_LIT:1> * ( n % block [ <NUM_LIT:0> ] != <NUM_LIT:0> ) ) <EOL> gridy = int ( m // block [ <NUM_LIT:1> ] + <NUM_LIT:1> * ( m % block [ <NUM_LIT:1> ] != <NUM_LIT:0> ) ) <EOL> grid = ( gridx , gridy , <NUM_LIT:1> ) <EOL> if out is None : <EOL> alloc = _global_cublas_allocator <EOL> out = gpuarray . empty_like ( x_gpu ) <EOL> else : <EOL> assert out . dtype == x_gpu . dtype <EOL> assert out . shape == x_gpu . shape <EOL> if x_gpu . flags . c_contiguous : <EOL> if axis == <NUM_LIT:0> : <EOL> col_kernel ( x_gpu , a_gpu , out , n , m , <EOL> block = block , grid = grid , stream = stream , <EOL> shared = <NUM_LIT> * x_gpu . dtype . itemsize ) <EOL> elif axis == <NUM_LIT:1> : <EOL> row_kernel ( x_gpu , a_gpu , out , n , m , <EOL> block = block , grid = grid , stream = stream , <EOL> shared = <NUM_LIT> * x_gpu . dtype . itemsize ) <EOL> else : <EOL> if axis == <NUM_LIT:0> : <EOL> row_kernel ( x_gpu , a_gpu , out , m , n , <EOL> block = block , grid = grid , stream = stream , <EOL> shared = <NUM_LIT> * x_gpu . dtype . itemsize ) <EOL> elif axis == <NUM_LIT:1> : <EOL> col_kernel ( x_gpu , a_gpu , out , m , n , <EOL> block = block , grid = grid , stream = stream , <EOL> shared = <NUM_LIT> * x_gpu . dtype . itemsize ) <EOL> return out <EOL> import operator <EOL> def binaryop_2d ( c_op , py_op , commutative , x_gpu , y_gpu ) : <EOL> if x_gpu . shape == y_gpu . shape : <EOL> return py_op ( x_gpu , y_gpu ) <EOL> elif x_gpu . size == <NUM_LIT:1> : <EOL> return py_op ( x_gpu . get ( ) . reshape ( ( ) ) , y_gpu ) <EOL> elif y_gpu . size == <NUM_LIT:1> : <EOL> return py_op ( x_gpu , y_gpu . get ( ) . reshape ( ( ) ) ) <EOL> if len ( x_gpu . shape ) == <NUM_LIT:2> : <EOL> m , n = x_gpu . shape <EOL> if y_gpu . shape == ( n , ) : <EOL> return binaryop_matvec ( c_op , x_gpu , y_gpu , axis = <NUM_LIT:1> ) <EOL> elif y_gpu . shape == ( <NUM_LIT:1> , n ) : <EOL> return binaryop_matvec ( c_op , x_gpu , y_gpu [ <NUM_LIT:0> ] , axis = <NUM_LIT:1> ) <EOL> elif y_gpu . shape == ( m , <NUM_LIT:1> ) : <EOL> return binaryop_matvec ( c_op , x_gpu , y_gpu . ravel ( ) , axis = <NUM_LIT:0> ) <EOL> if len ( y_gpu . shape ) == <NUM_LIT:2> and commutative : <EOL> m , n = y_gpu . shape <EOL> if x_gpu . shape == ( n , ) : <EOL> return binaryop_matvec ( c_op , y_gpu , x_gpu , axis = <NUM_LIT:1> ) <EOL> elif x_gpu . shape == ( <NUM_LIT:1> , n ) : <EOL> return binaryop_matvec ( c_op , y_gpu , x_gpu [ <NUM_LIT:0> ] , axis = <NUM_LIT:1> ) <EOL> elif x_gpu . shape == ( m , <NUM_LIT:1> ) : <EOL> return binaryop_matvec ( c_op , y_gpu , x_gpu . ravel ( ) , axis = <NUM_LIT:0> ) <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> def add ( x_gpu , y_gpu ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_2d ( "<STR_LIT:+>" , operator . add , True , x_gpu , y_gpu ) <EOL> def subtract ( x_gpu , y_gpu ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_2d ( "<STR_LIT:->" , operator . sub , False , x_gpu , y_gpu ) <EOL> def multiply ( x_gpu , y_gpu ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_2d ( "<STR_LIT:*>" , operator . mul , True , x_gpu , y_gpu ) <EOL> def divide ( x_gpu , y_gpu ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_2d ( "<STR_LIT:/>" , operator . truediv , False , x_gpu , y_gpu ) <EOL> def add_matvec ( x_gpu , a_gpu , axis = None , out = None , stream = None ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_matvec ( '<STR_LIT:+>' , x_gpu , a_gpu , axis , out , stream ) <EOL> def div_matvec ( x_gpu , a_gpu , axis = None , out = None , stream = None ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_matvec ( '<STR_LIT:/>' , x_gpu , a_gpu , axis , out , stream ) <EOL> def mult_matvec ( x_gpu , a_gpu , axis = None , out = None , stream = None ) : <EOL> """<STR_LIT>""" <EOL> return binaryop_matvec ( '<STR_LIT:*>' , x_gpu , a_gpu , axis , out , stream ) <EOL> def _sum_axis ( x_gpu , axis = None , out = None , calc_mean = False , ddof = <NUM_LIT:0> , <EOL> keepdims = False ) : <EOL> global _global_cublas_allocator <EOL> assert isinstance ( ddof , numbers . Integral ) <EOL> if axis is None or len ( x_gpu . shape ) <= <NUM_LIT:1> : <EOL> out_shape = ( <NUM_LIT:1> , ) * len ( x_gpu . shape ) if keepdims else ( ) <EOL> if calc_mean == False : <EOL> return gpuarray . sum ( x_gpu ) . reshape ( out_shape ) <EOL> else : <EOL> return gpuarray . sum ( x_gpu ) . reshape ( out_shape ) / ( x_gpu . dtype . type ( x_gpu . size - ddof ) ) <EOL> if axis < <NUM_LIT:0> : <EOL> axis += <NUM_LIT:2> <EOL> if axis > <NUM_LIT:1> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if x_gpu . flags . c_contiguous : <EOL> n , m = x_gpu . shape [ <NUM_LIT:1> ] , x_gpu . shape [ <NUM_LIT:0> ] <EOL> lda = x_gpu . shape [ <NUM_LIT:1> ] <EOL> trans = "<STR_LIT:n>" if axis == <NUM_LIT:0> else "<STR_LIT:t>" <EOL> sum_axis , out_axis = ( m , n ) if axis == <NUM_LIT:0> else ( n , m ) <EOL> else : <EOL> n , m = x_gpu . shape [ <NUM_LIT:0> ] , x_gpu . shape [ <NUM_LIT:1> ] <EOL> lda = x_gpu . shape [ <NUM_LIT:0> ] <EOL> trans = "<STR_LIT:t>" if axis == <NUM_LIT:0> else "<STR_LIT:n>" <EOL> sum_axis , out_axis = ( n , m ) if axis == <NUM_LIT:0> else ( m , n ) <EOL> if calc_mean : <EOL> alpha = ( <NUM_LIT:1.0> / ( sum_axis - ddof ) ) <EOL> else : <EOL> alpha = <NUM_LIT:1.0> <EOL> if ( x_gpu . dtype == np . complex64 ) : <EOL> gemv = cublas . cublasCgemv <EOL> elif ( x_gpu . dtype == np . float32 ) : <EOL> gemv = cublas . cublasSgemv <EOL> elif ( x_gpu . dtype == np . complex128 ) : <EOL> gemv = cublas . cublasZgemv <EOL> elif ( x_gpu . dtype == np . float64 ) : <EOL> gemv = cublas . cublasDgemv <EOL> alloc = _global_cublas_allocator <EOL> ons = ones ( ( sum_axis , ) , x_gpu . dtype , alloc ) <EOL> if keepdims : <EOL> out_shape = ( <NUM_LIT:1> , out_axis ) if axis == <NUM_LIT:0> else ( out_axis , <NUM_LIT:1> ) <EOL> else : <EOL> out_shape = ( out_axis , ) <EOL> if out is None : <EOL> out = gpuarray . empty ( out_shape , x_gpu . dtype , alloc ) <EOL> else : <EOL> assert out . dtype == x_gpu . dtype <EOL> assert out . size >= out_axis <EOL> gemv ( _global_cublas_handle , trans , n , m , <EOL> alpha , x_gpu . gpudata , lda , <EOL> ons . gpudata , <NUM_LIT:1> , <NUM_LIT:0.0> , out . gpudata , <NUM_LIT:1> ) <EOL> return out <EOL> def sum ( x_gpu , axis = None , out = None , keepdims = False ) : <EOL> """<STR_LIT>""" <EOL> return _sum_axis ( x_gpu , axis , out = out , keepdims = keepdims ) <EOL> def mean ( x_gpu , axis = None , out = None , keepdims = False ) : <EOL> """<STR_LIT>""" <EOL> return _sum_axis ( x_gpu , axis , calc_mean = True , out = out , keepdims = keepdims ) <EOL> def var ( x_gpu , ddof = <NUM_LIT:0> , axis = None , stream = None , keepdims = False ) : <EOL> """<STR_LIT>""" <EOL> def _inplace_pow ( x_gpu , p , stream ) : <EOL> func = elementwise . get_pow_kernel ( x_gpu . dtype ) <EOL> func . prepared_async_call ( x_gpu . _grid , x_gpu . _block , stream , <EOL> p , x_gpu . gpudata , x_gpu . gpudata , x_gpu . mem_size ) <EOL> if axis is None : <EOL> m = mean ( x_gpu ) . get ( ) <EOL> out = x_gpu - m <EOL> out **= <NUM_LIT:2> <EOL> out = _sum_axis ( out , axis = None , calc_mean = True , <EOL> ddof = ddof , out = None , keepdims = keepdims ) <EOL> else : <EOL> if axis < <NUM_LIT:0> : <EOL> axis += <NUM_LIT:2> <EOL> m = mean ( x_gpu , axis = axis ) <EOL> out = add_matvec ( x_gpu , - m , axis = <NUM_LIT:1> - axis , stream = stream ) <EOL> _inplace_pow ( out , <NUM_LIT:2> , stream ) <EOL> out = _sum_axis ( out , axis = axis , calc_mean = True , <EOL> ddof = ddof , out = None , keepdims = keepdims ) <EOL> return out <EOL> def std ( x_gpu , ddof = <NUM_LIT:0> , axis = None , stream = None , keepdims = False ) : <EOL> """<STR_LIT>""" <EOL> def _inplace_pow ( x_gpu , p , stream ) : <EOL> func = elementwise . get_pow_kernel ( x_gpu . dtype ) <EOL> func . prepared_async_call ( x_gpu . _grid , x_gpu . _block , stream , <EOL> p , x_gpu . gpudata , x_gpu . gpudata , x_gpu . mem_size ) <EOL> if axis is None : <EOL> return var ( x_gpu , ddof = ddof , stream = stream , keepdims = keepdims ) ** <NUM_LIT:0.5> <EOL> else : <EOL> out = var ( x_gpu , ddof = ddof , axis = axis , stream = stream , keepdims = keepdims ) <EOL> _inplace_pow ( out , <NUM_LIT:0.5> , stream ) <EOL> return out <EOL> @ context_dependent_memoize <EOL> def _get_minmax_kernel ( dtype , min_or_max ) : <EOL> template = Template ( """<STR_LIT>""" ) <EOL> cache_dir = None <EOL> ctype = dtype_to_ctype ( dtype ) <EOL> if min_or_max == '<STR_LIT>' : <EOL> iv = str ( np . finfo ( dtype ) . min ) <EOL> tmpl = template . substitute ( type = ctype , cmp_op = '<STR_LIT:>>' , init_value = iv ) <EOL> elif min_or_max == '<STR_LIT>' : <EOL> iv = str ( np . finfo ( dtype ) . max ) <EOL> tmpl = template . substitute ( type = ctype , cmp_op = '<STR_LIT:<>' , init_value = iv ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> mod = SourceModule ( tmpl ) <EOL> minmax_col_kernel = mod . get_function ( '<STR_LIT>' ) <EOL> minmax_row_kernel = mod . get_function ( '<STR_LIT>' ) <EOL> return minmax_col_kernel , minmax_row_kernel <EOL> def _minmax_impl ( a_gpu , axis , min_or_max , stream = None , keepdims = False ) : <EOL> '''<STR_LIT>''' <EOL> assert len ( a_gpu . shape ) < <NUM_LIT:3> <EOL> if iscomplextype ( a_gpu . dtype ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if axis is None or len ( a_gpu . shape ) <= <NUM_LIT:1> : <EOL> out_shape = ( <NUM_LIT:1> , ) * len ( a_gpu . shape ) <EOL> if min_or_max == '<STR_LIT>' : <EOL> return gpuarray . max ( a_gpu ) . reshape ( out_shape ) , None <EOL> else : <EOL> return gpuarray . min ( a_gpu ) . reshape ( out_shape ) , None <EOL> else : <EOL> if axis < <NUM_LIT:0> : <EOL> axis += <NUM_LIT:2> <EOL> assert axis in ( <NUM_LIT:0> , <NUM_LIT:1> ) <EOL> global _global_cublas_allocator <EOL> alloc = _global_cublas_allocator <EOL> n , m = a_gpu . shape if a_gpu . flags . c_contiguous else ( a_gpu . shape [ <NUM_LIT:1> ] , a_gpu . shape [ <NUM_LIT:0> ] ) <EOL> col_kernel , row_kernel = _get_minmax_kernel ( a_gpu . dtype , min_or_max ) <EOL> if ( axis == <NUM_LIT:0> and a_gpu . flags . c_contiguous ) or ( axis == <NUM_LIT:1> and a_gpu . flags . f_contiguous ) : <EOL> if keepdims : <EOL> out_shape = ( <NUM_LIT:1> , m ) if axis == <NUM_LIT:0> else ( m , <NUM_LIT:1> ) <EOL> else : <EOL> out_shape = ( m , ) <EOL> target = gpuarray . empty ( out_shape , dtype = a_gpu . dtype , allocator = alloc ) <EOL> idx = gpuarray . empty ( out_shape , dtype = np . uint32 , allocator = alloc ) <EOL> col_kernel ( a_gpu , target , idx , np . uint32 ( m ) , np . uint32 ( n ) , <EOL> block = ( <NUM_LIT:32> , <NUM_LIT:1> , <NUM_LIT:1> ) , grid = ( m , <NUM_LIT:1> , <NUM_LIT:1> ) , stream = stream ) <EOL> else : <EOL> if keepdims : <EOL> out_shape = ( <NUM_LIT:1> , n ) if axis == <NUM_LIT:0> else ( n , <NUM_LIT:1> ) <EOL> else : <EOL> out_shape = ( n , ) <EOL> target = gpuarray . empty ( out_shape , dtype = a_gpu , allocator = alloc ) <EOL> idx = gpuarray . empty ( out_shape , dtype = np . uint32 , allocator = alloc ) <EOL> row_kernel ( a_gpu , target , idx , np . uint32 ( m ) , np . uint32 ( n ) , <EOL> block = ( <NUM_LIT:32> , <NUM_LIT:1> , <NUM_LIT:1> ) , grid = ( n , <NUM_LIT:1> , <NUM_LIT:1> ) , stream = stream ) <EOL> return target , idx <EOL> def max ( a_gpu , axis = None , keepdims = False ) : <EOL> '''<STR_LIT>''' <EOL> return _minmax_impl ( a_gpu , axis , "<STR_LIT>" , keepdims = keepdims ) [ <NUM_LIT:0> ] <EOL> def min ( a_gpu , axis = None , keepdims = False ) : <EOL> '''<STR_LIT>''' <EOL> return _minmax_impl ( a_gpu , axis , "<STR_LIT>" , keepdims = keepdims ) [ <NUM_LIT:0> ] <EOL> def argmax ( a_gpu , axis , keepdims = False ) : <EOL> '''<STR_LIT>''' <EOL> if axis is None : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> return _minmax_impl ( a_gpu , axis , "<STR_LIT>" , keepdims = keepdims ) [ <NUM_LIT:1> ] <EOL> def argmin ( a_gpu , axis , keepdims = False ) : <EOL> '''<STR_LIT>''' <EOL> if axis is None : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> return _minmax_impl ( a_gpu , axis , "<STR_LIT>" , keepdims = keepdims ) [ <NUM_LIT:1> ] <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> import doctest <EOL> doctest . testmod ( ) </s>
