Conditional Assignment of Tensors

C++
1

Suppose that I have a two element-wise functions h(x) and g(x) where x is an at::Tensor.
How should I compute an element-wise conditional assignment; e.g.,
f(x) = h(x) if x > 0 else g(x)

2

Would at::where work for you? It’s the aten equivalent of torch.where.

Something like that:

at::Tensor f(const at::Tensor& x) {
    auto mask = x > 0;
    return at::where(mask, h(x), g(x));
}

The only issue is that h(x) and g(x) are computed for the entire x tensor, so they need to be defined for both positive and negative values. Also, if they’re slow to compute there could be faster alternatives that don’t require computing values that you won’t use. But I think in most cases where is good enough.

3

Yes, where is exactly what I was looking for. I need something that works when the non-selected components are nan. The following program demonstates that where works for this case:

#include <torch/torch.h>
int main() {
    using at::Tensor;
    //
    Tensor x = torch::tensor( {-4.0, 4.0} );
    Tensor g    = x * x;
    Tensor h    = x.sqrt();
    Tensor mask = x > 0;
    Tensor f    = at::where(mask, h, g);
    std::cout << f  << "\n";
}

I get the following output when I run that program above:

16
2
[ CPUFloatType{2} ]