How to implement numpy's np.lexsort in torch?

Viewed 150

I've checked that sorting a single column is faster on a GPU with torch than on a CPU with numpy, but I have a large Boolean array (~10^5 x ~10^3), and I need to sort the rows by the first column, breaking ties by the second column, breaking ties by the third, and so on. The numpy function I would use is np.lexsort. How do I do this using torch on a GPU?

By the way, there are more than 64 columns, or else I could turn each row into a single int64 and sort those.

0 Answers
Related