<s> import numpy <EOL> from chainer import cuda <EOL> from chainer import function <EOL> from chainer . utils import type_check <EOL> class Where ( function . Function ) : <EOL> """<STR_LIT>""" <EOL> def check_type_forward ( self , in_types ) : <EOL> type_check . expect ( in_types . size ( ) == <NUM_LIT:3> ) <EOL> c_type , x_type , y_type = in_types <EOL> type_check . expect ( <EOL> c_type . dtype == numpy . bool_ , <EOL> x_type . dtype == y_type . dtype , <EOL> x_type . shape == c_type . shape , <EOL> y_type . shape == c_type . shape , <EOL> ) <EOL> def forward ( self , inputs ) : <EOL> xp = cuda . get_array_module ( * inputs ) <EOL> condition , x , y = inputs <EOL> return xp . where ( condition , x , y ) , <EOL> def backward ( self , inputs , grads ) : <EOL> xp = cuda . get_array_module ( * inputs ) <EOL> condition , x , y = inputs <EOL> gx = xp . where ( condition , grads [ <NUM_LIT:0> ] , <NUM_LIT:0> ) <EOL> gy = xp . where ( condition , <NUM_LIT:0> , grads [ <NUM_LIT:0> ] ) <EOL> return None , gx , gy <EOL> def where ( condition , x , y ) : <EOL> """<STR_LIT>""" <EOL> return Where ( ) ( condition , x , y ) </s>
