Reset values less than threshold to 0 inplace

Hi!

Is it possible in vanilla eager PyTorch (without Inductor) to express the following inplace without allocating the output of x < 0.5?:

x = torch.empty(1000).uniform_(-10, 10)
y = torch.where(x < -1.5, 0, x, out = x)

torch.clamp, torch.hardswish, torch.hardtanh don’t cut it in an easy way

import torch

x = torch.empty(1000).uniform_(-10, 10)
print(x.data_ptr())  

x += 1.5
print(x.data_ptr())  # Should be the same if operation was in-place

x.clamp_(0)
print(x.data_ptr())  

x -= 1.5
print(x.data_ptr())  

x=x+1.5
print(x.data_ptr())

Would this work?

94261703466240
94261703466240
94261703466240
94261703466240
94261703474752

I don’t think this is correct computation. E.g. value -5 will be flushed to -1.5 and not to 0 after x -= 1.5

Oh, you are right. But right now I don’t think it’s possible. It’s a piecewise function and even if you don’t store the states, any math op needs to internally allocate some memory.
So you need a vectorized operator that really does if x<-1.5

Maybe you can achieve it with functorch.vmap — functorch nightly documentation
Otherwise you can write a quick cuda kernel.

Hi Vadim!

Use:

torch.nn.functional.threshold (x, -1.5, 0.0, inplace = True)

Best.

K. Frank