Question

我的问题非常相似 Indexing tensor with index matrix in theano? 除了我有3个维度。起初我想让它在numpy中工作。 2维度没有问题：

>>> idx = np.random.randint(3, size=(4, 2, 3))
>>> d = np.random.rand(4*2*3).reshape((4, 2, 3))
>>> d[1]
array([[ 0.37057415,  0.73066383,  0.76399376],
       [ 0.12155831,  0.12552545,  0.87648523]])
>>> idx[1]
array([[2, 0, 1],
       [2, 2, 2]])
>>> d[1][np.arange(d.shape[1])[:, np.newaxis], idx[1]]
array([[ 0.76399376,  0.37057415,  0.73066383],
       [ 0.87648523,  0.87648523,  0.87648523]])  #All correct

但我不知道如何让它适用于所有3个维度。尝试失败的示例：

>>> d[np.arange(d.shape[0])[:, np.newaxis], np.arange(d.shape[1]), idx]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (4,1) (2,) (4,2,3)

Answer 1

这有用吗？

d[
    np.arange(d.shape[0])[:, np.newaxis, np.newaxis],
    np.arange(d.shape[1])[:, np.newaxis],
    idx
]

您需要索引数组共同拥有可广播的维度

theano中的高级3d indexind

1 个答案: