Tensorflow：每行索引

Question

假設我有一個形狀的張量（100,20）。現在我也有一個形狀指數的張量（100，）。現在如何獲得形狀（100）或（100,1），每行（100行）權值（由相應的指數指標選擇的張量

小例子：？ 所以我們可以說張量是

[1, 2, 3] 
[4, 5, 6] 
[7, 8, 9] 

和張量B是

[0,2,1] 

然後我想作爲輸出

[1,6,8] 

Answer 1

你可以加入你的乙張量的適當範圍，以創建二維指數（在你的例子[[0, 0], [1, 2], [2, 1]]），然後使用tf.gather_nd提取元素：

b_2 = tf.expand_dims(b, 1) 
range = tf.expand_dims(tf.range(tf.shape(b)[0]), 1) 
ind = tf.concat(1, [range, b_2]) 
res = tf.gather_nd(a, ind)