Numpy: сортировка многомерного массива по многомерному массиву

Простите меня, если это излишне или слишком просто. Я перехожу на Python/Numpy из R и с трудом переворачиваю вещи в голове.

У меня есть n-мерный массив, который я хочу отсортировать, используя другой n-мерный массив значений индекса. Я знаю, что мог бы зациклить это, но кажется, что должен быть действительно краткий Numpyonic способ победить это в представлении. Вот мой пример кода для решения задачи, где n=2:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 

Итак, теперь у меня есть массив случайных чисел 2 x 5 и индекс 2 x 5. Я читал справку по take()примерно 10 раз, но мой мозг, очевидно, не вникает в нее.

Я думал, что это может привести меня туда:

take(a1, index)

array([[ 0.29589188, -0.71279375, -0.18154864, -1.12184984,  0.25698875],
       [ 0.29589188, -0.71279375, -0.18154864,  0.25698875, -1.12184984]])

но это явно переупорядочивает только первый элемент (я полагаю, из-за выравнивания).

Любые советы о том, как мне перейти от того места, где я нахожусь, к решению, которое сортирует элемент 0 из a1 по элементу 0 индекса ... элемента n?