theano 中的高级 3d indexind

Question

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

>>> 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
]

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