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 

給我需要的輸出（我想這個象徵性的版本）。

Similiarly在Theano，

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？

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

不像NumPy的，的index的長度（=的input第一尺寸）是不固定的例子。

Answer 1

你可以做切片與tf.gather_nd：

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