Actually there’s an issue: I find the sort order of identical items to be non-deterministic (probably on GPU only). For example on the row (part of a 2-D matrix)
[1023,    8, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023,
606, 1023, 1023, 1023, 1023, 1023,  408, 1023, 1023, 1023,  427, 1023]
I got the result
[ 408,  427,    8,  606, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023,
1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023, 1023]
(1023 acts as 0 here).