Tensorflow中的高级索引

Question

我有两个形状

的张量

voxels :: (4, 64, 64, 64) indices :: (4096, 3)

此处voxels有dtype float32，而indices是0和63之间的整数。 我想索引体素，创建一个张量张量output：

output :: 4, 4096

以下约束成立

output[i,j] = voxels[i, indices[j,0], indices[j,1]. indices[j,2]]

如何在tensorflow中执行此操作？

有一些关于更高级切片的讨论here，但不清楚那些建议是否适用，或者渐变已经实施。

作为参考，如果上面的张量是numpy数组，下面的代码就是我想要的：

output = voxels[:, indices[:,0],indices[:,1],indices[:,2]]

Answer 1

gather_nd确实是一种可能性，但渐变尚未实现。但是，也可以使用涉及gather的技巧，通过展平张量并重新计算适当的指数。

Answer 2

我相信gather_nd会做你想要的，虽然我不认为渐变已经实现。