Sort a subpart of an Aten tensor?

(Will Shen) #1

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.

#2

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.

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;
}
1 Like
(Will Shen) #3

Thanks for the help. Appreciate it.