<s> import cupy <EOL> def roll ( a , shift , axis = None ) : <EOL> """<STR_LIT>""" <EOL> if axis is None : <EOL> if a . size == <NUM_LIT:0> : <EOL> return a <EOL> size = a . size <EOL> ra = a . ravel ( ) <EOL> shift %= size <EOL> res = cupy . empty ( ( size , ) , a . dtype ) <EOL> res [ : shift ] = ra [ size - shift : ] <EOL> res [ shift : ] = ra [ : size - shift ] <EOL> return res . reshape ( a . shape ) <EOL> else : <EOL> axis = int ( axis ) <EOL> if axis < <NUM_LIT:0> : <EOL> axis += a . ndim <EOL> if not <NUM_LIT:0> <= axis < a . ndim : <EOL> raise ValueError ( '<STR_LIT>' % ( - a . ndim , a . ndim ) ) <EOL> size = a . shape [ axis ] <EOL> if size == <NUM_LIT:0> : <EOL> return a <EOL> shift %= size <EOL> prev = ( slice ( None ) , ) * axis <EOL> rest = ( slice ( None ) , ) * ( a . ndim - axis - <NUM_LIT:1> ) <EOL> ind1 = prev + ( slice ( size - shift , None , None ) , ) + rest <EOL> ind2 = prev + ( slice ( None , size - shift , None ) , ) + rest <EOL> r_ind1 = prev + ( slice ( None , shift , None ) , ) + rest <EOL> r_ind2 = prev + ( slice ( shift , None , None ) , ) + rest <EOL> res = cupy . empty_like ( a ) <EOL> res [ r_ind1 ] = a [ ind1 ] <EOL> res [ r_ind2 ] = a [ ind2 ] <EOL> return res </s>
