Creating BoolTensor Mask in torch C++

Viewed 593

I am trying to create a mask for torch in C++ of type BoolTensor. The first n elements in dimension one need to be False and the rest need to be True.

This is my attempt but I do not know if this is correct (size is the number of elements):

src_mask = torch::BoolTensor({6, 1});
src_mask[:size,:] = 0;
src_mask[size:,:] = 1;
2 Answers

I'm not sure to understand exactly your goal here, so here is my best attempt to convert into C++ you pseudo-code .

First, with libtorch you declare the type of your tensor through the torch::TensorOptions struct (types names are prefixed with a lowercase k)

Second, your python-like slicing is possible thanks to the torch::Tensor::slicefunction (see here and there).

Finally, that gives you something like :

// Creates a tensor of boolean, initially all ones
auto options = torch::TensorOptions().dtype(torch::kBool));
torch::Tensor bool_tensor = torch::ones({6,1}, options);
// Set the slice to 0
int size = 3;
bool_tensor.slice(/*dim=*/0, /*start=*/0, /*end=*/size) = 0;

std::cout << bool_tensor << std::endl;

Please not that this will set the first size rows to 0. I assumed that's what you meant by "first elements in dimension x".

Another way to do it:

using namespace torch::indexing; //for using Slice(...) function

at::Tensor src_mask = at::empty({ 6, 1 }, at::kBool); //empty bool tensor 
src_mask.index_put_({ Slice(None, size), Slice() }, 0); //src_mask[:size,:] = 0
src_mask.index_put_({ Slice(size, None), Slice() }, 1); //src_mask[size:,:] = 0
Related