在 MXNet NDArray 的列上使用索引数组

Question

给定一个索引数组 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() and 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)>

最后，要解决这个问题，有没有办法就地执行此操作？

Update 这是一个稍微更优雅但可能效率不高（由于换位）的上述版本：

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

Answer 1

您需要的是 MXNet 中所谓的高级索引。提交了一个 PR，用于通过 MXNet NDArray 的高级索引获取元素，并且还将向 NDArray 添加设置元素的功能。预计1.0版本出来。

Using Index Arrays on Columns of an MXNet NDArray