動的インデックスのリストを集める

Question

埋め込みセットを取得する方法を知る手助けは誰ですか？

私は、各インデックスの確率を予測するいくつかのコードを持って、その後、最大の選択：

# U is batch_size x max_sentence_length x embedding_size 
scores_per_index = find_start_preds(U ...) # batch_size x max_sentence_length x 1 
start_preds = tf.argmax(alpha, axis=1) # batch_size x 1 

それぞれに関連する可能性、単語を再つかむ埋め込みが予測を開始する場合、私は、したいと思います。それは可能ですか？これは私が考えているものですが、それは動作しません:(

u_s = U[:, start_preds, :] 

Answer 1

あなたが望むもののためにtf.gatherを使用することができるはずですが、あなたは順序を変更する必要があるので、それが唯一の主要な指標に動作します：

U2 = tf.transpose(U, [1, 0, 2]) 
u_s = tf.gather(U2, start_preds) 
u_s = tf.transpose(u_s, [1, 0, 2])