[ATen] C++ API: Equivalent to Conditional Index

Hi everyone,

I am using ATen’s tensor library in my C++ project, and I am having a problem now.

In Python, if I have a tensor tmp, then I can do the following code to add 100 to all the values in tmp that is less than zero.

tmp[tmp < 0].add_(100)

I hope to do the similar in C++. I try to do:

tmp.index(tmp < 0).add_(100)

However, it seems index() create a new tensor storage and the change only happens on new tensor, the original tmp is unchanged.

I checked other similar functions, e.g. index_select, and they all do not meet the need.

Would anyone know this part and please give me some advice.

Thank you!

Another way to do this in python is
tmp + (tmp < 0).type_as(tmp) * 100

Does something like that work for you in ATen?

1 Like

Fantastic idea! Thank you, and I think it is something I can start with.