使用张量索引

Question

我试图使用Targmax张量来索引张量。

在numpy中，您可以执行以下索引：

mat = np.random.uniform(size = 3*10*10).reshape((3,10,10))
indices = [np.array([0,0,1,2]),np.array([1,1,2,3]), np.array([1,3,0,3])]
mat[indices]

tensorflow中是否有等效的操作？

Answer 1

x = tf.constant([[1,2],[3,4]])
sess = tf.Session()
sess.run(tf.gather_nd(x,[[0,0],[1,1]]))

停止

array([1, 4], dtype=int32)