<s> import numpy as np <EOL> from numba import types <EOL> from numba . extending import overload <EOL> @ overload ( np . where ) <EOL> def where ( cond , x , y ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( cond , types . Array ) : <EOL> if all ( ty . layout == '<STR_LIT:C>' for ty in ( cond , x , y ) ) : <EOL> def where_impl ( cond , x , y ) : <EOL> """<STR_LIT>""" <EOL> shape = cond . shape <EOL> if x . shape != shape or y . shape != shape : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> res = np . empty_like ( x ) <EOL> cf = cond . flat <EOL> xf = x . flat <EOL> yf = y . flat <EOL> rf = res . flat <EOL> for i in range ( cond . size ) : <EOL> rf [ i ] = xf [ i ] if cf [ i ] else yf [ i ] <EOL> return res <EOL> else : <EOL> def where_impl ( cond , x , y ) : <EOL> """<STR_LIT>""" <EOL> shape = cond . shape <EOL> if x . shape != shape or y . shape != shape : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> res = np . empty_like ( x ) <EOL> for idx , c in np . ndenumerate ( cond ) : <EOL> res [ idx ] = x [ idx ] if c else y [ idx ] <EOL> return res <EOL> else : <EOL> def where_impl ( cond , x , y ) : <EOL> """<STR_LIT>""" <EOL> scal = x if cond else y <EOL> return np . full_like ( scal , scal ) <EOL> return where_impl </s>
