Hi,
I am learning the source code of pytorch, and I would like to know how can pytorch compute the rank along cerain dimension of a tensor. However, the code is too abstract to understand, so I am asking what method is used to compute this please ? Did you use thrust combined with a for loop to compute it many times, or you wrote a cuda kernel with merge sort or methods like this?
Hi,
argssort is implemented by doing sort(same_args)[1] 
You can see that implementation here: pytorch/LegacyDefinitions.cpp at 22a34bcf4e5eaa348f0117c414c3dd760ec64b13 · pytorch/pytorch · GitHub
And for sort, you can see here that it comes from the old THC library: pytorch/native_functions.yaml at master · pytorch/pytorch · GitHub
And you can find the THC implementation in a couple files there: pytorch/THCTensorSort.cu at 070a30b265c1f4f404081ffefa7c41383750cc2e · pytorch/pytorch · GitHub