TensorFlowのインデキシングサブセンサー

Question

入力テンソルのindex〜サブセンサーを適用して出力テンソルを取得したいとします。 numpyの、

import numpy as np 

input = np.random.random((100,5)) # matrix 
index = np.randint(5, size=(100,)) # vector 
output = data[np.arange(index.shape[0]), index] # vector 

で例えば

は（私はこれを象徴するバージョンが必要）私に所望の出力を提供します。 TheanoでSimiliarly

、

import theano.tensor as T 
import theano 

input = T.matrix() # symbolic matrix 
index = T.ivector() # symbolic vector 
output = input[T.arange(index.shape[0]), index] # symbolic vector 

は私に希望outputを与えます。

どうすればいいですか？TensorFlow？ numpyの、index（inputの= 1番目の寸法）の長さを有する例が固定されていないとは異なり

import tensorflow as tf 
input = tf.placeholder('float32', [None, 5]) 
index = tf.placeholder('int32', [None]) 
output = ??? 

。

Answer 1

あなたはtf.gather_ndとスライスを行うことができます。

output = tf.gather_nd(input, tf.stack((tf.range(tf.shape(index)[0]), index), -1))