Использование индексных массивов в столбцах MXNet NDArray

Учитывая массив индексов index и, скажем, матрицу A, мне нужна матрица B с соответствующей перестановкой столбцов A.

В Numpy я бы сделал следующее:

>>> A = np.arange(6).reshape(2,3); A
array([[0, 1, 2],
       [3, 4, 5]])
>>> index = [2,0,1]
>>> A[:,index]
array([[2, 0, 1],
       [5, 3, 4]])

Есть ли естественный или эффективный способ сделать это в MXNet? Функции pick() и take(), похоже, не работает таким образом. Мне удалось придумать следующее, но это не элегантно.

>>> mx.nd.take(A.T, mx.nd.array([[2],[0],[1]])).T.reshape((2,3))

[[ 2.  0.  1.]
 [ 5.  3.  4.]]
<NDArray 2x3 @cpu(0)>

Наконец, чтобы бросить гаечный ключ в работу, есть ли способ сделать это на месте?

Обновление Вот немного более элегантная, но предположительно не такая эффективная (из-за транспонирования) версия выше:

>>> mx.nd.take(A.T, mx.nd.array([2,0,1])).T
[[ 2.  0.  1.]
 [ 5.  3.  4.]]
<NDArray 2x3 @cpu(0)>