<s> """<STR_LIT>""" <EOL> from . library import * <EOL> from . util import * <EOL> from . util import _is_number <EOL> from . base import * <EOL> from . bcast import _bcast_var <EOL> import math <EOL> class Seq ( ct . Structure ) : <EOL> """<STR_LIT>""" <EOL> _fields_ = [ ( "<STR_LIT>" , ct . c_double ) , <EOL> ( "<STR_LIT:end>" , ct . c_double ) , <EOL> ( "<STR_LIT>" , ct . c_double ) ] <EOL> def __init__ ( self , S ) : <EOL> self . begin = ct . c_double ( <NUM_LIT:0> ) <EOL> self . end = ct . c_double ( - <NUM_LIT:1> ) <EOL> self . step = ct . c_double ( <NUM_LIT:1> ) <EOL> if _is_number ( S ) : <EOL> self . begin = ct . c_double ( S ) <EOL> self . end = ct . c_double ( S ) <EOL> elif isinstance ( S , slice ) : <EOL> if ( S . step is not None ) : <EOL> self . step = ct . c_double ( S . step ) <EOL> if ( S . step < <NUM_LIT:0> ) : <EOL> self . begin , self . end = self . end , self . begin <EOL> if ( S . start is not None ) : <EOL> self . begin = ct . c_double ( S . start ) <EOL> if ( S . stop is not None ) : <EOL> self . end = ct . c_double ( S . stop - math . copysign ( <NUM_LIT:1> , self . step ) ) <EOL> else : <EOL> raise IndexError ( "<STR_LIT>" ) <EOL> class ParallelRange ( Seq ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , start , stop = None , step = None ) : <EOL> if ( stop is None ) : <EOL> stop = start <EOL> start = <NUM_LIT:0> <EOL> self . S = slice ( start , stop , step ) <EOL> super ( ParallelRange , self ) . __init__ ( self . S ) <EOL> def __iter__ ( self ) : <EOL> return self <EOL> def next ( self ) : <EOL> """<STR_LIT>""" <EOL> if _bcast_var . get ( ) is True : <EOL> _bcast_var . toggle ( ) <EOL> raise StopIteration <EOL> else : <EOL> _bcast_var . toggle ( ) <EOL> return self <EOL> def __next__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . next ( ) <EOL> class _uidx ( ct . Union ) : <EOL> _fields_ = [ ( "<STR_LIT>" , ct . c_void_p ) , <EOL> ( "<STR_LIT>" , Seq ) ] <EOL> class Index ( ct . Structure ) : <EOL> _fields_ = [ ( "<STR_LIT>" , _uidx ) , <EOL> ( "<STR_LIT>" , ct . c_bool ) , <EOL> ( "<STR_LIT>" , ct . c_bool ) ] <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , idx ) : <EOL> self . idx = _uidx ( ) <EOL> self . isBatch = False <EOL> self . isSeq = True <EOL> if isinstance ( idx , BaseArray ) : <EOL> arr = ct . c_void_p ( <NUM_LIT:0> ) <EOL> if ( idx . type ( ) == Dtype . b8 . value ) : <EOL> safe_call ( backend . get ( ) . af_where ( ct . pointer ( arr ) , idx . arr ) ) <EOL> else : <EOL> safe_call ( backend . get ( ) . af_retain_array ( ct . pointer ( arr ) , idx . arr ) ) <EOL> self . idx . arr = arr <EOL> self . isSeq = False <EOL> elif isinstance ( idx , ParallelRange ) : <EOL> self . idx . seq = idx <EOL> self . isBatch = True <EOL> else : <EOL> self . idx . seq = Seq ( idx ) <EOL> def __del__ ( self ) : <EOL> if not self . isSeq : <EOL> arr = ct . c_void_p ( self . idx . arr ) <EOL> backend . get ( ) . af_release_array ( arr ) <EOL> class _Index4 ( object ) : <EOL> def __init__ ( self , idx0 , idx1 , idx2 , idx3 ) : <EOL> index_vec = Index * <NUM_LIT:4> <EOL> self . array = index_vec ( idx0 , idx1 , idx2 , idx3 ) <EOL> self . idxs = [ idx0 , idx1 , idx2 , idx3 ] <EOL> @ property <EOL> def pointer ( self ) : <EOL> return ct . pointer ( self . array ) <EOL> def __getitem__ ( self , idx ) : <EOL> return self . array [ idx ] <EOL> def __setitem__ ( self , idx , value ) : <EOL> self . array [ idx ] = value <EOL> self . idxs [ idx ] = value </s>
