How can I do the operation the same as `np.where`?
wangg12
DiffEverything
jekbradbury
I actually solved it by a workaround: <SCODE># differentiable equivalent of np.where
# cond could be a FloatTensor with zeros and ones
def where(cond, x_1, x_2):
    return (cond * x_1) + ((1-cond) * x_2)
<ECODE>
jaromiru
yxchng
truenicoco
thanks others for examples of other options for 0.3 and before
Pete_Florence
<SCODE>## this example prunes any elements in the vector above or below the bounds
vec = torch.FloatTensor([-3.2, 1000.0, 10.0, 639.0])
lower_bound = 0.0
upper_bound = 640.0
lower_bound_vec = torch.ones_like(vec) * lower_bound
upper_bound_vec = torch.ones_like(vec) * upper_bound
zeros_vec       = torch.zeros_like(vec)

def where(cond, x_1, x_2):
    cond = cond.float()    
    return (cond * x_1) + ((1-cond) * x_2)

vec = where(vec < lower_bound_vec, zeros_vec, vec)
vec = where(vec > upper_bound_vec, zeros_vec, vec)

in_bound_indices = torch.nonzero(vec).squeeze(1)
vec = torch.index_select(vec, 0, in_bound_indices)
<ECODE>
Pete_Florence
jpeg729
Pete_Florence