Как вернуть индексы из сортировки двухмерного массива строк по строкам?

Question

Ввод: двумерный массив NumPy

Вывод: массив индексов, который будет сортировать массив по строкам (или по столбцам)

Например: скажем, функция get_sorted_indices (массив, ось = 0)

a = np.array([[1,2,3,4,5]
             ,[2,3,4,5,6]
             ,[1,2,3,4,5]
             ,[2,3,4,6,6]
             ,[2,3,4,5,6]])

ind = get_sorted_indices(a, axis=0)

Тогда мы получим

>>> ind 
[0, 2, 1, 4, 3]
>>> a[ind]   # should be equals to a.sort(axis = 0)
array([[1, 2, 3, 4, 5],
       [1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [2, 3, 4, 5, 6],
       [2, 3, 4, 6, 6]])
>>> a.sort(axis=0) 
array([[1, 2, 3, 4, 5],
       [1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [2, 3, 4, 5, 6],
       [2, 3, 4, 6, 6]])

Я посмотрел на argsort, но не понимаю его вывод, и чтение документации не поможет:

>>> a.argsort() 
array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])

>>> a.argsort(axis=0) 
array([[0, 0, 0, 0, 0],
       [2, 2, 2, 2, 2],
       [1, 1, 1, 1, 1],
       [3, 3, 3, 4, 3],
       [4, 4, 4, 3, 4]])

Я могу сделать это вручную, но мне кажется, что я неправильно понимаю argsort или я что-то упускаю из numpy. Есть ли стандартный способ сделать это, или у меня нет выбора, кроме как сделать это вручную?