Hi there,
I’m wondering if we can sort a subpart of an ATen tensor, specified with a start index and an end index? More generally, can we create inplace an ATen sub-tensor from full tensor? Similar to C++ array, where we just pass in a different pointer.
Thanks.
You can simply use std::sort (for CPU tensors) or thrust::sort (for CUDA tensors) on the tensor’s data pointer, given that your tensor is contiguous.
std::sort
thrust::sort
The snippet below is built for libtorch, but the behavior is exactly the same in ATen:
#include <torch/torch.h> #include <iostream> int main(){ torch::Tensor x = torch::tensor({1, 3, 4, 2, 0, 5, 9}); int *x_ptr = x.data<int>(); std::sort(x_ptr + 1, x_ptr + 4, std::less<int>()); // sorting {3, 4, 2} std::cout << x << std::endl; // prints 1 2 3 4 0 5 9 return 0; }
Thanks for the help. Appreciate it.