Ordenar invariante para numpy.argsort con múltiples dimensiones

numpy.argsort docs state

Devoluciones:
index_array: ndarray, int Matriz de índices que ordenan a lo largo del eje especificado. Si a es unidimensional,a[index_array] produce un ordenado a.

¿Cómo puedo aplicar el resultado denumpy.argsort para que una matriz multidimensional recupere una matriz ordenada? (NO solo una matriz 1-D o 2-D; podría ser una matriz N-dimensional donde N solo se conoce en tiempo de ejecución)

>>> import numpy as np
>>> np.random.seed(123)
>>> A = np.random.randn(3,2)
>>> A
array([[-1.0856306 ,  0.99734545],
       [ 0.2829785 , -1.50629471],
       [-0.57860025,  1.65143654]])
>>> i=np.argsort(A,axis=-1)
>>> A[i]
array([[[-1.0856306 ,  0.99734545],
        [ 0.2829785 , -1.50629471]],

       [[ 0.2829785 , -1.50629471],
        [-1.0856306 ,  0.99734545]],

       [[-1.0856306 ,  0.99734545],
        [ 0.2829785 , -1.50629471]]])

Para mí no es solo cuestión de usarsort() en lugar; Tengo otra matrizB y quiero ordenarB utilizando los resultados denp.argsort(A) a lo largo del eje apropiado. Considere el siguiente ejemplo:

>>> A = np.array([[3,2,1],[4,0,6]])
>>> B = np.array([[3,1,4],[1,5,9]])
>>> i = np.argsort(A,axis=-1)
>>> BsortA = ???             
# should result in [[4,1,3],[5,1,9]]
# so that corresponding elements of B and sort(A) stay together

Parece que esta funcionalidad esya una solicitud de mejora en numpy.

Respuestas a la pregunta(3)

Su respuesta a la pregunta