This seems the right answer…
The first row is 3,0,1,2 since 0.2267 is smallest (which is at index 3), then 0.4143 (which is at index 0), then 0.4293, then 0.9506.
Similarly for second and third row.
Note argsort is done by default along the last dimension.
1 Like